quant-os/tests/test_pit_universe.py

import unittest
from datetime import datetime, timezone

from quant60.pit import (
    PITIntegrityError,
    PITRecord,
    PointInTimeTable,
    assert_information_available,
)
from quant60.universe import (
    SecurityEligibility,
    UniverseState,
    build_universe,
    classify_security,
)


UTC = timezone.utc


def record(
    *,
    value,
    effective="2024-03-31T00:00:00+08:00",
    available="2024-04-30T18:00:00+08:00",
    revision=0,
):
    return PITRecord(
        entity_id="600000.XSHG",
        field="net_profit",
        effective_time=effective,
        available_time=available,
        value=value,
        source="fixture",
        source_version="v1",
        revision=revision,
    )


class PITUniverseTests(unittest.TestCase):
    def test_pit_query_never_uses_future_announcement(self):
        table = PointInTimeTable([record(value=100)])
        before = table.query(
            "600000.XSHG",
            "net_profit",
            as_of="2024-04-30T17:59:59+08:00",
        )
        after = table.query(
            "600000.XSHG",
            "net_profit",
            as_of="2024-04-30T18:00:00+08:00",
        )
        self.assertIsNone(before)
        self.assertEqual(after.value, 100)

    def test_revision_is_visible_only_after_its_own_available_time(self):
        table = PointInTimeTable(
            [
                record(value=100),
                record(
                    value=90,
                    available="2024-05-05T09:00:00+08:00",
                    revision=1,
                ),
            ]
        )
        first = table.query(
            "600000.XSHG",
            "net_profit",
            as_of="2024-05-01T15:00:00+08:00",
        )
        revised = table.query(
            "600000.XSHG",
            "net_profit",
            as_of="2024-05-06T15:00:00+08:00",
        )
        self.assertEqual(first.value, 100)
        self.assertEqual(revised.value, 90)

    def test_future_injection_gate_fails_closed(self):
        future = record(
            value=100,
            available="2024-05-02T09:00:00+08:00",
        )
        with self.assertRaisesRegex(PITIntegrityError, "future-available"):
            assert_information_available(
                [future],
                as_of="2024-05-01T15:00:00+08:00",
            )

    def test_naive_timestamps_and_duplicate_identity_are_rejected(self):
        with self.assertRaisesRegex(PITIntegrityError, "timezone"):
            record(
                value=100,
                effective=datetime(2024, 3, 31),
            )
        item = record(value=100)
        with self.assertRaisesRegex(PITIntegrityError, "duplicate"):
            PointInTimeTable([item, item])

    def test_data_version_is_order_independent(self):
        left = record(value=100)
        right = PITRecord(
            entity_id="000001.XSHE",
            field="net_profit",
            effective_time=datetime(2024, 3, 31, tzinfo=UTC),
            available_time=datetime(2024, 4, 30, tzinfo=UTC),
            value=80,
            source="fixture",
            source_version="v1",
        )
        self.assertEqual(
            PointInTimeTable([left, right]).data_version,
            PointInTimeTable([right, left]).data_version,
        )

    def test_universe_states_keep_untradable_holdings_visible(self):
        frozen = classify_security(
            SecurityEligibility(
                symbol="600000.SH",
                index_member=False,
                listing_sessions=500,
                median_turnover=1_000_000,
                minimum_turnover=100_000,
                suspended=True,
                held_quantity=1000,
                sellable_quantity=1000,
            )
        )
        sell_only = classify_security(
            SecurityEligibility(
                symbol="000001.SZ",
                index_member=True,
                listing_sessions=500,
                median_turnover=1_000_000,
                minimum_turnover=100_000,
                is_st=True,
                held_quantity=1000,
                sellable_quantity=1000,
            )
        )
        self.assertEqual(frozen.state, UniverseState.FROZEN)
        self.assertTrue(frozen.can_hold)
        self.assertFalse(frozen.can_sell)
        self.assertEqual(sell_only.state, UniverseState.SELL_ONLY)
        self.assertTrue(sell_only.can_sell)
        hold_only = classify_security(
            SecurityEligibility(
                symbol="000333.SZ",
                index_member=True,
                listing_sessions=500,
                median_turnover=50_000,
                minimum_turnover=100_000,
                held_quantity=1000,
                sellable_quantity=1000,
            )
        )
        self.assertEqual(hold_only.state, UniverseState.HOLD_ONLY)
        self.assertTrue(hold_only.can_hold)
        self.assertFalse(hold_only.can_open)

    def test_openable_and_excluded_are_distinct(self):
        decisions = build_universe(
            [
                SecurityEligibility(
                    symbol="600000.SH",
                    index_member=True,
                    listing_sessions=500,
                    median_turnover=1_000_000,
                    minimum_turnover=100_000,
                ),
                SecurityEligibility(
                    symbol="300750.SZ",
                    index_member=True,
                    listing_sessions=20,
                    median_turnover=1_000_000,
                    minimum_turnover=100_000,
                ),
            ]
        )
        self.assertEqual(
            decisions["600000.XSHG"].state,
            UniverseState.OPENABLE,
        )
        self.assertEqual(
            decisions["300750.XSHE"].state,
            UniverseState.EXCLUDED,
        )


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