source/tests/test_publish_wrapup_evidence.py

import copy
import hashlib
import inspect
import json
from pathlib import Path
import subprocess
import tempfile
from types import SimpleNamespace
import unittest

from scripts import publish_wrapup_evidence as evidence


class FakeHub:
    def __init__(self, root, corrupt=False):
        self.root, self.corrupt = root, corrupt
        self.files = {'README.md': b'---\nlicense: unknown\n---\nFinal card\n',
                      'FINAL_MODEL.json': b'{"final":"preserved"}',
                      'CURRENT_SNAPSHOT.json': b'{"training":"preserved"}'}
        self.uploads = self.pointer_commits = self.downloads = 0

    def repo_info(self, *args, **kwargs):
        return SimpleNamespace(private=True, sha='current-head')

    def upload_folder(self, **kwargs):
        self.uploads += 1
        folder = Path(kwargs['folder_path'])
        for path in folder.rglob('*'):
            if path.is_file():
                self.files[kwargs['path_in_repo'] + '/' + str(path.relative_to(folder))] = path.read_bytes()
        return SimpleNamespace(oid='payload-commit')

    def get_paths_info(self, repo_id, paths, **kwargs):
        result = []
        for name in paths:
            if name not in self.files:
                continue
            content = self.files[name]
            checksum = hashlib.sha256(content).hexdigest()
            if self.corrupt and name.endswith('.pt'):
                checksum = '0' * 64
            result.append(SimpleNamespace(path=name, size=len(content), lfs={'sha256': checksum} if name.endswith('.pt') else None,
                blob_id=hashlib.sha1(f'blob {len(content)}\0'.encode() + content).hexdigest()))
        return result

    def hf_hub_download(self, repo_id, filename, **kwargs):
        if filename.endswith(('.pt', '.tar.gz')):
            raise AssertionError('Large weights/source must not be redownloaded')
        self.downloads += 1
        path = self.root / f'download-{self.downloads}'
        path.write_bytes(self.files[filename])
        return str(path)

    def create_commit(self, repo_id, operations, **kwargs):
        self.pointer_commits += 1
        assert [operation.path_in_repo for operation in operations] == ['PROFILE_RESULTS.json']
        for operation in operations:
            self.files[operation.path_in_repo] = operation.path_or_fileobj
        return SimpleNamespace(oid='pointer-commit')


class WrapupEvidenceTests(unittest.TestCase):
    def setUp(self):
        temporary = tempfile.TemporaryDirectory()
        self.addCleanup(temporary.cleanup)
        self.root = Path(temporary.name)
        self.source, self.inputs = self.root / 'source', self.root / 'inputs'
        self.source.mkdir()
        self.inputs.mkdir()
        (self.source / 'training.py').write_text('# committed reconstruction source\n')
        self.git('init', '-q')
        self.commit()
        self.revision = self.git('rev-parse', 'HEAD').strip()
        (self.inputs / 'checkpoint.pt').write_bytes(b'trusted fixture adapter and optimizer checkpoint')
        (self.inputs / 'summary.json').write_text('{"status":"complete"}\n')
        (self.inputs / 'proof.json').write_text('{"cpu_verified":true}\n')
        (self.inputs / 'report.md').write_text('Frozen profiling report.\n')
        self.plan = {'format': 'opensysone-wrapup-evidence-plan-v1', 'repo_id': 'andyshu/opensysone',
            'evidence_id': '20260917T120000Z-wrapup', 'files': {
                'artifacts/gx10/checkpoint.pt': {'source': str(self.inputs / 'checkpoint.pt')},
                'profiling/summary.json': {'source': str(self.inputs / 'summary.json')},
                'proofs/checkpoint.json': {'source': str(self.inputs / 'proof.json')},
                'reports/report.md': {'source': str(self.inputs / 'report.md')}},
            'completion_checks': [{'path': 'profiling/summary.json', 'json_equals': {'status': 'complete'}}],
            'checkpoints': [{'path': 'artifacts/gx10/checkpoint.pt', 'sha256': evidence.digest(self.inputs / 'checkpoint.pt'),
                'step': 159, 'source_commit': self.revision, 'role': 'resumable', 'proof': 'proofs/checkpoint.json'}],
            'provenance': {'selection_frozen_before_profiling': True}}
        self.plan_path, self.output = self.root / 'plan.json', self.root / 'exports/evidence'

    def git(self, *args):
        return subprocess.check_output(['git', *args], cwd=self.source, text=True, stderr=subprocess.DEVNULL)

    def commit(self):
        self.git('add', '.')
        self.git('-c', 'user.name=Fixture', '-c', 'user.email=fixture@example.invalid', 'commit', '-qm', 'fixture')

    def prepare(self):
        evidence.write_json(self.plan_path, self.plan)
        return evidence.prepare(self.plan_path, self.output, self.source, (self.root,))

    def test_complete_payload_updates_only_supplemental_pointer_with_installed_hub_signatures(self):
        from huggingface_hub import CommitOperationAdd, HfApi
        from huggingface_hub.hf_api import RepoFile
        manifest = self.prepare()
        api = FakeHub(self.root)
        originals = dict(api.files)
        for name in ('repo_info', 'upload_folder', 'get_paths_info', 'hf_hub_download', 'create_commit'):
            method, signature = getattr(api, name), inspect.signature(getattr(HfApi, name))
            def checked(*args, _method=method, _signature=signature, _name=name, **kwargs):
                _signature.bind(api, *args, **kwargs)
                value = _method(*args, **kwargs)
                if _name == 'get_paths_info':
                    value = [RepoFile(path=item.path, size=item.size, oid=item.blob_id,
                        lfs={'size': item.size, 'oid': item.lfs['sha256'], 'pointerSize': 128} if item.lfs else None) for item in value]
                return value
            setattr(api, name, checked)
        result = evidence.publish(api, 'andyshu/opensysone', self.output, lambda *a, **k: None,
                                  operation_factory=CommitOperationAdd)
        self.assertEqual(result['pointer_commit'], 'pointer-commit')
        self.assertEqual(manifest['source_revisions'], [self.revision])
        self.assertIn(f"profiles/{self.plan['evidence_id']}/payload/sources/{self.revision}/source.tar.gz", api.files)
        for name, content in originals.items():
            self.assertEqual(api.files[name], content)

    def test_publish_accepts_actual_progress_callback_signature_and_preserves_phase(self):
        self.prepare()
        state, phases = {}, []
        # Match main's callback signature and update behavior, rather than a
        # variadic no-op that would hide duplicate positional/keyword arguments.
        def progress(stage, **values):
            state.update(stage=stage, heartbeat_utc=evidence.utc(), **values)
            phases.append(state['stage'])
            evidence.write_json(self.root / 'publication-state.json', state)
        evidence.publish(FakeHub(self.root), 'andyshu/opensysone', self.output, progress,
                         operation_factory=SimpleNamespace)
        saved = evidence.read_json(self.root / 'publication-state.json')
        self.assertEqual(phases, ['uploading_payload', 'verifying_payload', 'publishing_evidence_pointer'])
        self.assertEqual(saved['stage'], 'publishing_evidence_pointer')
        self.assertEqual(saved['export_directory'], str(self.output))

    def test_incomplete_evidence_and_bad_checkpoint_hash_block_staging(self):
        (self.inputs / 'summary.json').write_text('{"status":"running"}')
        with self.assertRaisesRegex(ValueError, 'not complete'):
            self.prepare()
        self.assertFalse(self.output.exists())
        (self.inputs / 'summary.json').write_text('{"status":"complete"}')
        self.plan['checkpoints'][0]['sha256'] = '0' * 64
        with self.assertRaisesRegex(ValueError, 'expected checksum'):
            self.prepare()
        self.assertFalse(self.output.exists())

    def test_protocol_checksum_and_exact_exit_code_preserved_with_completion_gate(self):
        checksum, exit_code = self.inputs / 'protocol.sha256', self.inputs / 'exit_code'
        checksum.write_text('a' * 64 + '\n')
        exit_code.write_text('0\n')
        self.plan['files'].update({'profiling/prepared/protocol.sha256': {'source': str(checksum)},
                                   'profiling/wrapper/exit_code': {'source': str(exit_code)}})
        self.plan['completion_checks'].append({'path': 'profiling/wrapper/exit_code', 'text_equals': '0'})
        manifest = self.prepare()
        for name, source in [('profiling/prepared/protocol.sha256', checksum), ('profiling/wrapper/exit_code', exit_code)]:
            self.assertEqual((self.output / 'payload' / name).read_bytes(), source.read_bytes())
            self.assertEqual(manifest['files']['payload/' + name]['sha256'], evidence.digest(source))

    def test_other_extensionless_sources_and_renamed_exit_code_rejected(self):
        original = copy.deepcopy(self.plan)
        for name in ('token', 'noextensionsecret', 'Exit_Code'):
            with self.subTest(source=name):
                self.plan = copy.deepcopy(original)
                path = self.inputs / name
                path.write_text('never publish')
                self.plan['files']['proofs/exit_code'] = {'source': str(path)}
                with self.assertRaisesRegex(ValueError, 'prohibited'):
                    self.prepare()
        self.plan = copy.deepcopy(original)
        path = self.inputs / 'exit_code'
        path.write_text('0\n')
        self.plan['files']['proofs/other_extensionless'] = {'source': str(path)}
        with self.assertRaisesRegex(ValueError, 'disguise'):
            self.prepare()

    def test_traversal_secret_and_base_weight_sources_are_blocked(self):
        original = copy.deepcopy(self.plan)
        for name in ('../FINAL_MODEL.json', '/absolute.json'):
            with self.subTest(name=name):
                self.plan = copy.deepcopy(original)
                self.plan['files'][name] = {'source': str(self.inputs / 'proof.json')}
                with self.assertRaises(ValueError):
                    self.prepare()
        for name in ('credentials.json', '.env', 'api-key.json', 'model.safetensors'):
            with self.subTest(name=name):
                self.plan = copy.deepcopy(original)
                path = self.inputs / name
                path.write_text('never publish')
                self.plan['files']['proofs/extra' + path.suffix] = {'source': str(path)}
                with self.assertRaises(ValueError):
                    self.prepare()
        self.plan = copy.deepcopy(original)
        models = self.root / 'models'
        models.mkdir()
        (models / 'best.pt').write_bytes(b'base weights disguised as adapter')
        self.plan['files']['artifacts/gx10/checkpoint.pt']['source'] = str(models / 'best.pt')
        with self.assertRaisesRegex(ValueError, 'base-model directories'):
            self.prepare()

    def test_source_archive_containing_weights_is_blocked(self):
        (self.source / 'model.safetensors').write_bytes(b'forbidden source weight')
        self.commit()
        with self.assertRaisesRegex(ValueError, 'prohibited'):
            self.prepare()
        self.assertFalse(self.output.exists())

    def test_remote_mismatch_never_updates_any_pointer(self):
        self.prepare()
        api = FakeHub(self.root, corrupt=True)
        originals = dict(api.files)
        with self.assertRaisesRegex(ValueError, 'Remote evidence checksum'):
            evidence.publish(api, 'andyshu/opensysone', self.output, lambda *a, **k: None)
        self.assertEqual(api.pointer_commits, 0)
        self.assertNotIn('PROFILE_RESULTS.json', api.files)
        for name, content in originals.items():
            self.assertEqual(api.files[name], content)

    def test_corrupt_stage_blocks_upload_and_existing_id_is_immutable(self):
        self.prepare()
        api = FakeHub(self.root)
        path = self.output / 'payload/reports/report.md'
        path.chmod(0o644)
        original = path.read_bytes()
        path.write_bytes(b'changed')
        with self.assertRaisesRegex(ValueError, 'stage checksum'):
            evidence.publish(api, 'andyshu/opensysone', self.output, lambda *a, **k: None)
        self.assertEqual(api.uploads, 0)
        path.write_bytes(original)
        api.files[f"profiles/{self.plan['evidence_id']}/manifest.json"] = b'other immutable content'
        with self.assertRaisesRegex(ValueError, 'different content'):
            evidence.publish(api, 'andyshu/opensysone', self.output, lambda *a, **k: None)
        self.assertEqual(api.uploads, 0)


if __name__ == '__main__':
    unittest.main()