tests/test_integration.py
1.7 KB · 57 lines · python Raw
1 """End-to-end integration tests."""
2
3 from __future__ import annotations
4
5 from pqc_hypervisor_attestation import (
6 AttestationVerifier,
7 Attester,
8 ContinuousAttester,
9 InMemoryBackend,
10 RegionSnapshot,
11 )
12
13 WORKLOAD_ID = "model-serving-1"
14
15
16 def test_full_flow_register_attest_verify(
17 attester: Attester, backend: InMemoryBackend
18 ) -> None:
19 loop = ContinuousAttester(
20 attester=attester, backend=backend, workload_id=WORKLOAD_ID
21 )
22 # Pin the expected hashes for both regions.
23 loop.expected_hashes = {
24 "model-weights-0": RegionSnapshot.hash_bytes(b"\xaa" * 128),
25 "activation-cache": RegionSnapshot.hash_bytes(b"\xbb" * 64),
26 }
27 report = loop.attest_once()
28 result = AttestationVerifier.verify(report, strict=True)
29 assert result.valid is True
30 assert result.drifts == []
31
32
33 def test_tampered_region_flagged_after_attestation(
34 attester: Attester, backend: InMemoryBackend
35 ) -> None:
36 loop = ContinuousAttester(
37 attester=attester, backend=backend, workload_id=WORKLOAD_ID
38 )
39 loop.expected_hashes = {
40 "model-weights-0": RegionSnapshot.hash_bytes(b"\xaa" * 128),
41 "activation-cache": RegionSnapshot.hash_bytes(b"\xbb" * 64),
42 }
43
44 # First attestation: clean.
45 ok_report = loop.attest_once()
46 assert AttestationVerifier.verify(ok_report, strict=True).valid is True
47
48 # Simulated attacker mutates weights in place.
49 backend.update("model-weights-0", b"\xff" * 128)
50
51 # Next attestation: drift detected via expected_hash vs new snapshot.
52 drift_report = loop.attest_once()
53 result = AttestationVerifier.verify(drift_report, strict=True)
54 assert result.signature_valid is True
55 assert "model-weights-0" in result.drifts
56 assert result.valid is False
57