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()