siloqy/prod/tests/test_dita_v2_kernel.py

from __future__ import annotations

from datetime import datetime, timezone
import unittest

from prod.clean_arch.dita_v2 import (
    AccountProjection,
    BackendMode,
    ControlUpdate,
    ExecutionKernel,
    InMemoryControlPlane,
    InMemoryZincPlane,
    KernelCommandType,
    KernelControlSnapshot,
    KernelEventKind,
    KernelIntent,
    KernelMode,
    KernelVerbosity,
    MemoryKernelJournal,
    MockVenueAdapter,
    MockVenueScenario,
    TradeSide,
    TradeSlot,
    TradeStage,
    VenueEvent,
    VenueEventStatus,
)


def mk_intent(
    *,
    action: KernelCommandType = KernelCommandType.ENTER,
    slot_id: int = 0,
    trade_id: str = "trade-1",
    asset: str = "BTCUSDT",
    side: TradeSide = TradeSide.SHORT,
    target_size: float = 1.0,
    leverage: float = 2.0,
    reference_price: float = 100.0,
    exit_leg_ratios=(1.0,),
    reason: str = "TEST",
) -> KernelIntent:
    return KernelIntent(
        timestamp=datetime.now(timezone.utc),
        intent_id=f"intent-{trade_id}-{action.value}",
        trade_id=trade_id,
        slot_id=slot_id,
        asset=asset,
        side=side,
        action=action,
        reference_price=reference_price,
        target_size=target_size,
        leverage=leverage,
        exit_leg_ratios=tuple(exit_leg_ratios),
        reason=reason,
    )


class TestDITAv2ControlPlane(unittest.TestCase):
    def test_control_plane_updates_and_mirrors(self):
        plane = InMemoryControlPlane()
        updated = plane.update(
            ControlUpdate(
                mode=KernelMode.DEBUG,
                verbosity=KernelVerbosity.TRACE,
                backend_mode=BackendMode.BINGX,
                trace_transitions=True,
            )
        )
        self.assertEqual(updated.mode, KernelMode.DEBUG)
        self.assertEqual(updated.verbosity, KernelVerbosity.TRACE)
        self.assertEqual(updated.backend_mode, BackendMode.BINGX)
        self.assertTrue(updated.trace_transitions)
        self.assertEqual(plane.mirror()["mode"], KernelMode.DEBUG.value)


class TestDITAv2Kernel(unittest.TestCase):
    def test_entry_ack_fill_reaches_position_open(self):
        journal = MemoryKernelJournal()
        zinc = InMemoryZincPlane()
        venue = MockVenueAdapter(MockVenueScenario(emit_fill_on_submit=True, partial_fill_ratio=1.0))
        kernel = ExecutionKernel(
            control_plane=InMemoryControlPlane(
                KernelControlSnapshot(mode=KernelMode.DEBUG, verbosity=KernelVerbosity.TRACE)
            ),
            venue=venue,
            journal=journal,
            zinc_plane=zinc,
        )

        outcome = kernel.process_intent(mk_intent())

        slot = kernel.slot(0)
        self.assertTrue(outcome.accepted)
        self.assertEqual(slot.fsm_state, TradeStage.POSITION_OPEN)
        self.assertFalse(slot.closed)
        self.assertEqual(slot.trade_id, "trade-1")
        self.assertAlmostEqual(slot.size, 1.0, places=6)
        self.assertEqual(len(journal.rows), 3)
        self.assertEqual(len(zinc.intent_region), 1)
        self.assertEqual(zinc.read_control().mode, KernelMode.DEBUG)

    def test_partial_fill_stays_working_then_full_fill_opens_position(self):
        journal = MemoryKernelJournal()
        venue = MockVenueAdapter(MockVenueScenario(emit_fill_on_submit=False, partial_fill_ratio=0.5))
        kernel = ExecutionKernel(
            control_plane=InMemoryControlPlane(
                KernelControlSnapshot(mode=KernelMode.DEBUG, verbosity=KernelVerbosity.TRACE)
            ),
            venue=venue,
            journal=journal,
        )

        kernel.process_intent(mk_intent())
        slot = kernel.slot(0)
        self.assertEqual(slot.fsm_state, TradeStage.ENTRY_WORKING)
        self.assertAlmostEqual(slot.size, 0.5, places=6)

        full_fill = VenueEvent(
            timestamp=datetime.now(timezone.utc),
            event_id="evt-full",
            trade_id="trade-1",
            slot_id=0,
            kind=KernelEventKind.FULL_FILL,
            status=VenueEventStatus.FILLED,
            venue_order_id=slot.active_entry_order.venue_order_id if slot.active_entry_order else "V-00000001",
            venue_client_id=slot.active_entry_order.venue_client_id if slot.active_entry_order else "trade-1:intent-trade-1-ENTER",
            side=TradeSide.SHORT,
            asset="BTCUSDT",
            price=100.0,
            size=1.0,
            filled_size=1.0,
            remaining_size=0.0,
        )
        kernel.on_venue_event(full_fill)

        self.assertEqual(slot.fsm_state, TradeStage.POSITION_OPEN)
        self.assertFalse(slot.closed)
        self.assertAlmostEqual(slot.size, 1.0, places=6)

    def test_two_leg_exit_closes_only_after_final_leg(self):
        kernel = ExecutionKernel(
            control_plane=InMemoryControlPlane(
                KernelControlSnapshot(mode=KernelMode.DEBUG, verbosity=KernelVerbosity.TRACE)
            ),
            venue=MockVenueAdapter(MockVenueScenario(emit_fill_on_submit=True, partial_fill_ratio=1.0)),
            journal=MemoryKernelJournal(),
        )

        kernel.process_intent(mk_intent())
        slot = kernel.slot(0)
        slot.exit_leg_ratios = (0.5, 0.5)

        first_exit = kernel.process_intent(
            mk_intent(action=KernelCommandType.EXIT, target_size=0.5, exit_leg_ratios=(0.5, 0.5), reason="TP1")
        )
        self.assertTrue(first_exit.accepted)
        self.assertEqual(slot.fsm_state, TradeStage.POSITION_OPEN)
        self.assertFalse(slot.closed)
        self.assertAlmostEqual(slot.size, 0.5, places=6)

        second_exit = kernel.process_intent(
            mk_intent(action=KernelCommandType.EXIT, target_size=0.5, exit_leg_ratios=(0.5, 0.5), reason="TP2")
        )
        self.assertTrue(second_exit.accepted)
        self.assertTrue(slot.closed)
        self.assertEqual(slot.fsm_state, TradeStage.CLOSED)
        self.assertAlmostEqual(slot.size, 0.0, places=6)

    def test_reconcile_sets_stale_state(self):
        kernel = ExecutionKernel(
            control_plane=InMemoryControlPlane(),
            venue=MockVenueAdapter(),
            journal=MemoryKernelJournal(),
        )
        kernel.process_intent(mk_intent())
        slot = kernel.slot(0)
        kernel.process_intent(mk_intent(action=KernelCommandType.RECONCILE))
        self.assertEqual(slot.fsm_state, TradeStage.STALE_STATE_RECONCILING)

    def test_account_projection_aggregates_slots(self):
        projection = AccountProjection()
        slots = [
            TradeSlot(
                slot_id=0,
                trade_id="t1",
                asset="BTCUSDT",
                side=TradeSide.SHORT,
                entry_price=100.0,
                size=1.0,
                initial_size=1.0,
                leverage=2.0,
                fsm_state=TradeStage.POSITION_OPEN,
                metadata={"mark_price": 99.0},
            ),
            TradeSlot(
                slot_id=1,
                trade_id="t2",
                asset="ETHUSDT",
                side=TradeSide.LONG,
                entry_price=50.0,
                size=2.0,
                initial_size=2.0,
                leverage=3.0,
                fsm_state=TradeStage.EXIT_WORKING,
                metadata={"mark_price": 55.0},
            ),
        ]

        projection.observe_slots(slots)
        self.assertEqual(projection.snapshot.open_positions, 2)
        self.assertAlmostEqual(projection.snapshot.open_notional, 209.0, places=6)
        self.assertGreater(projection.snapshot.leverage, 0.0)

    def test_debug_mode_journal_records_transitions(self):
        journal = MemoryKernelJournal()
        kernel = ExecutionKernel(
            control_plane=InMemoryControlPlane(
                KernelControlSnapshot(mode=KernelMode.DEBUG, verbosity=KernelVerbosity.TRACE)
            ),
            venue=MockVenueAdapter(MockVenueScenario(emit_fill_on_submit=True, partial_fill_ratio=1.0)),
            journal=journal,
        )

        kernel.process_intent(mk_intent())
        self.assertGreaterEqual(len(journal.rows), 2)
        self.assertTrue(all("slot_state" in row for row in journal.rows))


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