tests/conftest.py
2.5 KB · 89 lines · python Raw
1 """Shared fixtures for pqc-mbom tests."""
2
3 from __future__ import annotations
4
5 import hashlib
6
7 import pytest
8 from quantumshield.identity.agent import AgentIdentity
9
10 from pqc_mbom import (
11 ComponentType,
12 LicenseInfo,
13 MBOM,
14 MBOMBuilder,
15 ModelComponent,
16 )
17
18
19 @pytest.fixture
20 def creator_identity() -> AgentIdentity:
21 return AgentIdentity.create("test-mbom-creator")
22
23
24 @pytest.fixture
25 def base_arch_component() -> ModelComponent:
26 return ModelComponent(
27 component_id="base-llama3-0001",
28 component_type=ComponentType.BASE_ARCHITECTURE,
29 name="Llama-3",
30 version="3.0",
31 content_hash=hashlib.sha3_256(b"llama-3-architecture").hexdigest(),
32 supplier="Meta",
33 license=LicenseInfo(spdx_id="llama-3-community", name="Llama 3 Community License"),
34 )
35
36
37 @pytest.fixture
38 def weights_component() -> ModelComponent:
39 return ModelComponent(
40 component_id="weights-0001",
41 component_type=ComponentType.WEIGHTS,
42 name="llama3-8b.safetensors",
43 content_hash=hashlib.sha3_256(b"fake-weights-blob").hexdigest(),
44 content_size=16_000_000_000,
45 supplier="Meta",
46 )
47
48
49 @pytest.fixture
50 def training_data_component() -> ModelComponent:
51 return ModelComponent(
52 component_id="train-cc-2024",
53 component_type=ComponentType.TRAINING_DATA,
54 name="common-crawl-2024",
55 content_hash=hashlib.sha3_256(b"cc-dataset-manifest").hexdigest(),
56 content_size=1_000_000_000_000,
57 supplier="Common Crawl",
58 external_url="https://commoncrawl.org/",
59 )
60
61
62 @pytest.fixture
63 def tokenizer_component() -> ModelComponent:
64 return ModelComponent(
65 component_id="tok-llama3",
66 component_type=ComponentType.TOKENIZER,
67 name="llama3-tokenizer",
68 content_hash=hashlib.sha3_256(b"tokenizer-vocab").hexdigest(),
69 )
70
71
72 @pytest.fixture
73 def sample_components(
74 base_arch_component: ModelComponent,
75 weights_component: ModelComponent,
76 training_data_component: ModelComponent,
77 tokenizer_component: ModelComponent,
78 ) -> list[ModelComponent]:
79 return [base_arch_component, weights_component, training_data_component, tokenizer_component]
80
81
82 @pytest.fixture
83 def sample_mbom(sample_components: list[ModelComponent]) -> MBOM:
84 builder = MBOMBuilder("Llama-3-8B-Instruct", "1.0.0", supplier="Meta")
85 builder.set_description("Test MBOM for unit tests")
86 for c in sample_components:
87 builder.add_component(c)
88 return builder.build()
89