policyengine-us-data/policyengine_us_data/tests/conftest.py at 32f0d639d887b73876cef082df4d21d5b778cbd7 · PolicyEngine/policyengine-us-data · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
"""Shared fixtures for version manifest tests."""

from unittest.mock import MagicMock

import pytest

from policyengine_us_data.utils.version_manifest import (
    HFVersionInfo,
    GCSVersionInfo,
    VersionManifest,
    VersionRegistry,
)


@pytest.fixture
def sample_generations() -> dict[str, int]:
    return {
        "enhanced_cps_2024.h5": 1710203948123456,
        "cps_2024.h5": 1710203948234567,
        "states/AL.h5": 1710203948345678,
    }


@pytest.fixture
def sample_hf_info() -> HFVersionInfo:
    return HFVersionInfo(
        repo="policyengine/policyengine-us-data",
        commit="abc123def456",
    )


@pytest.fixture
def sample_manifest(
    sample_generations: dict[str, int],
    sample_hf_info: HFVersionInfo,
) -> VersionManifest:
    return VersionManifest(
        version="1.72.3",
        created_at="2026-03-10T14:30:00Z",
        hf=sample_hf_info,
        gcs=GCSVersionInfo(
            bucket="policyengine-us-data",
            generations=sample_generations,
        ),
    )


@pytest.fixture
def sample_registry(
    sample_manifest: VersionManifest,
) -> VersionRegistry:
    """A registry with one version entry."""
    return VersionRegistry(
        current="1.72.3",
        versions=[sample_manifest],
    )


@pytest.fixture
def mock_bucket() -> MagicMock:
    bucket = MagicMock()
    bucket.name = "policyengine-us-data"
    return bucket