tests/conftest.py
1.1 KB · 47 lines · python Raw
1 """Shared pytest fixtures."""
2
3 from __future__ import annotations
4
5 import pytest
6 from quantumshield.identity.agent import AgentIdentity
7
8 from pqc_hypervisor_attestation import (
9 Attester,
10 InMemoryBackend,
11 MemoryRegion,
12 )
13
14 WORKLOAD_ID = "model-serving-1"
15
16
17 @pytest.fixture
18 def attester_identity() -> AgentIdentity:
19 return AgentIdentity.create("test-attester", capabilities=["attest"])
20
21
22 @pytest.fixture
23 def attester(attester_identity: AgentIdentity) -> Attester:
24 return Attester(attester_identity)
25
26
27 @pytest.fixture
28 def backend() -> InMemoryBackend:
29 be = InMemoryBackend()
30 weights = MemoryRegion(
31 region_id="model-weights-0",
32 description="Llama weights shard 0",
33 address=0x1000,
34 size=128,
35 protection="RO",
36 )
37 cache = MemoryRegion(
38 region_id="activation-cache",
39 description="KV cache for in-flight request",
40 address=0x2000,
41 size=64,
42 protection="RW",
43 )
44 be.register(WORKLOAD_ID, weights, b"\xaa" * 128)
45 be.register(WORKLOAD_ID, cache, b"\xbb" * 64)
46 return be
47