"""Add new tests to test_contract_uploads.py."""
import pathlib

p = pathlib.Path('tests/test_contract_uploads.py')
t = p.read_bytes()

# 1. Header-only test in TestSignedUpload (before 'class TestTheHelper:')
marker = b'\n\nclass TestTheHelper:'
test = (
    b"\n    def test_a_header_only_pdf_is_refused(\n"
    b"        self, app, client, as_role, make_user, contract_for\n"
    b"    ):\n"
    b'        """A file with a %PDF- header but no real PDF structure is rejected."""\n'
    b"        player_id = as_role('player')\n"
    b"        admin_id = make_user('admin')\n"
    b"        contract_id, directory = contract_for(player_id, admin_id)\n\n"
    b"        client.post(\n"
    b"            f'/users/contracts/{contract_id}/upload_signed',\n"
    b"            data={'signed_file': (io.BytesIO(PDF_HEADER_ONLY_BYTES), 'signed.pdf')},\n"
    b"            content_type='multipart/form-data',\n"
    b"            follow_redirects=True,\n"
    b"        )\n\n"
    b"        with app.app_context():\n"
    b"            contract = db.session.get(Contract, contract_id)\n"
    b"            assert contract.status != 'signed'\n"
    b"            assert contract.signed_file_path is None\n\n"
)
pos = t.find(marker)
if pos >= 0:
    t = t[:pos] + test + t[pos:]

# 2. Header-only helper test at end of TestTheHelper
trailing = t.rstrip(b'\n')
idx = t.find(trailing) + len(trailing)
helper = (
    b"\n    def test_a_header_only_file_is_refused(self, app):\n"
    b'        """A stream starting with %PDF- but holding non-PDF data is\n'
    b"        rejected (the header check is not enough).\"\"\"\n"
    b"        from werkzeug.datastructures import FileStorage\n\n"
    b"        from app.routes.users import ALLOWED_CONTRACT_EXTENSIONS, pdf_upload_error\n\n"
    b"        upload = FileStorage(stream=io.BytesIO(PDF_HEADER_ONLY_BYTES), filename='clever.pdf')\n"
    b"        with app.test_request_context('/'):\n"
    b"            assert pdf_upload_error(upload, ALLOWED_CONTRACT_EXTENSIONS)\n"
)
t = t[:idx] + helper + t[idx:]

p.write_bytes(t)
print('Added new tests')