test_engine.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. from __future__ import annotations
  2. from pathlib import Path
  3. from tempfile import TemporaryDirectory
  4. from fastapi.testclient import TestClient
  5. from src.trader_mcp import strategy_registry, strategy_store, strategy_engine
  6. from src.trader_mcp.server import app
  7. STRATEGY_CODE = '''
  8. from src.trader_mcp.strategy_sdk import Strategy
  9. class Strategy(Strategy):
  10. CONFIG_SCHEMA = {"label": {"type": "string", "default": "hello world"}}
  11. def init(self):
  12. return {"started": True, "config_copy": dict(self.config)}
  13. def on_tick(self, tick):
  14. self.state["ticks"] = self.state.get("ticks", 0) + 1
  15. return self.state["ticks"]
  16. def render(self):
  17. return {"widgets": [{"type": "metric", "label": "ticks", "value": self.state.get("ticks", 0)}]}
  18. '''
  19. def test_mode_off_does_not_instantiate_and_active_does(tmp_path):
  20. original_db = strategy_store.DB_PATH
  21. original_dir = strategy_registry.STRATEGIES_DIR
  22. try:
  23. strategy_store.DB_PATH = tmp_path / "trader_mcp.sqlite3"
  24. strategy_registry.STRATEGIES_DIR = tmp_path / "strategies"
  25. strategy_registry.STRATEGIES_DIR.mkdir()
  26. (strategy_registry.STRATEGIES_DIR / "demo.py").write_text(STRATEGY_CODE)
  27. strategy_store.add_strategy_instance(id="s1", strategy_type="demo", account_id="acct-1", client_id="cid-1", mode="off", config={"x": 1})
  28. result = strategy_engine.reconcile_all()
  29. assert result["running"] == []
  30. assert strategy_engine.get_running_strategy("s1") is None
  31. strategy_store.update_strategy_mode("s1", "active")
  32. result = strategy_engine.reconcile_all()
  33. assert "s1" in result["running"]
  34. runtime = strategy_engine.get_running_strategy("s1")
  35. assert runtime is not None
  36. assert runtime.instance.state["started"] is True
  37. assert runtime.instance.context.account_id == "acct-1"
  38. assert runtime.instance.context.client_id == "cid-1"
  39. tick_result = strategy_engine.tick_strategy("s1", {"price": 1})
  40. assert tick_result["ok"] is True
  41. assert tick_result["result"] == 1
  42. render_result = strategy_engine.render_strategy("s1")
  43. assert render_result["ok"] is True
  44. assert render_result["render"]["widgets"][0]["value"] == 1
  45. finally:
  46. strategy_store.DB_PATH = original_db
  47. strategy_registry.STRATEGIES_DIR = original_dir
  48. strategy_engine._running.clear()
  49. def test_mode_change_route_reconciles(tmp_path):
  50. original_db = strategy_store.DB_PATH
  51. original_dir = strategy_registry.STRATEGIES_DIR
  52. try:
  53. strategy_store.DB_PATH = tmp_path / "trader_mcp.sqlite3"
  54. strategy_registry.STRATEGIES_DIR = tmp_path / "strategies"
  55. strategy_registry.STRATEGIES_DIR.mkdir()
  56. (strategy_registry.STRATEGIES_DIR / "demo.py").write_text(STRATEGY_CODE)
  57. client = TestClient(app)
  58. client.post(
  59. "/strategies",
  60. json={"id": "s2", "strategy_type": "demo", "account_id": "acct-1", "client_id": "cid-1", "mode": "off", "config": {"x": 1}},
  61. )
  62. r = client.post("/strategies/s2/mode", json={"mode": "active"})
  63. assert r.status_code == 200
  64. assert "s2" in r.json()["running"]
  65. finally:
  66. strategy_store.DB_PATH = original_db
  67. strategy_registry.STRATEGIES_DIR = original_dir
  68. strategy_engine._running.clear()
  69. def test_dashboard_add_strategy_synthesizes_client_and_defaults(tmp_path):
  70. original_db = strategy_store.DB_PATH
  71. original_dir = strategy_registry.STRATEGIES_DIR
  72. try:
  73. strategy_store.DB_PATH = tmp_path / "trader_mcp.sqlite3"
  74. strategy_registry.STRATEGIES_DIR = tmp_path / "strategies"
  75. strategy_registry.STRATEGIES_DIR.mkdir()
  76. (strategy_registry.STRATEGIES_DIR / "hello_world.py").write_text(STRATEGY_CODE)
  77. client = TestClient(app)
  78. response = client.post(
  79. "/dashboard/strategies/add",
  80. data={"name": "My super Grid 0.5", "strategy_type": "hello_world", "account_id": "acct-1"},
  81. follow_redirects=False,
  82. )
  83. assert response.status_code in {302, 303}
  84. record = strategy_store.list_strategy_instances()[0]
  85. assert record.name == "My super Grid 0.5"
  86. assert record.client_id.startswith("hello_world:")
  87. assert record.config == {"label": "hello world"}
  88. assert record.mode == "off"
  89. finally:
  90. strategy_store.DB_PATH = original_db
  91. strategy_registry.STRATEGIES_DIR = original_dir
  92. strategy_engine._running.clear()
  93. def test_runtime_pause_suppresses_tick_and_render(tmp_path):
  94. original_db = strategy_store.DB_PATH
  95. original_dir = strategy_registry.STRATEGIES_DIR
  96. try:
  97. strategy_store.DB_PATH = tmp_path / "trader_mcp.sqlite3"
  98. strategy_registry.STRATEGIES_DIR = tmp_path / "strategies"
  99. strategy_registry.STRATEGIES_DIR.mkdir()
  100. (strategy_registry.STRATEGIES_DIR / "hello_world.py").write_text(STRATEGY_CODE)
  101. strategy_store.add_strategy_instance(id="s3", strategy_type="hello_world", account_id="acct-1", client_id="cid-1", mode="active", config={})
  102. strategy_engine.reconcile_all()
  103. assert strategy_engine.pause_strategy("s3")["ok"] is True
  104. tick_result = strategy_engine.tick_strategy("s3", {"price": 1})
  105. assert tick_result["paused"] is True
  106. assert tick_result["skipped"] is True
  107. render_result = strategy_engine.render_strategy("s3")
  108. assert render_result["paused"] is True
  109. assert render_result["render"] is None
  110. assert strategy_engine.resume_strategy("s3")["ok"] is True
  111. tick_result = strategy_engine.tick_strategy("s3", {"price": 1})
  112. assert tick_result["result"] == 1
  113. render_result = strategy_engine.render_strategy("s3")
  114. assert render_result["render"]["widgets"][0]["value"] == 1
  115. finally:
  116. strategy_store.DB_PATH = original_db
  117. strategy_registry.STRATEGIES_DIR = original_dir
  118. strategy_engine._running.clear()
  119. def test_dashboard_pause_toggle(tmp_path):
  120. original_db = strategy_store.DB_PATH
  121. original_dir = strategy_registry.STRATEGIES_DIR
  122. try:
  123. strategy_store.DB_PATH = tmp_path / "trader_mcp.sqlite3"
  124. strategy_registry.STRATEGIES_DIR = tmp_path / "strategies"
  125. strategy_registry.STRATEGIES_DIR.mkdir()
  126. (strategy_registry.STRATEGIES_DIR / "hello_world.py").write_text(STRATEGY_CODE)
  127. client = TestClient(app)
  128. client.post(
  129. "/dashboard/strategies/add",
  130. data={"name": "Pause test", "strategy_type": "hello_world", "account_id": "acct-1"},
  131. follow_redirects=False,
  132. )
  133. strategy_id = strategy_store.list_strategy_instances()[0].id
  134. strategy_store.update_strategy_mode(strategy_id, "active")
  135. strategy_engine.reconcile_instance(strategy_id)
  136. assert strategy_engine.get_running_strategy(strategy_id).paused is False
  137. client.post(f"/dashboard/strategies/{strategy_id}/pause", follow_redirects=False)
  138. assert strategy_engine.get_running_strategy(strategy_id).paused is True
  139. client.post(f"/dashboard/strategies/{strategy_id}/pause", follow_redirects=False)
  140. assert strategy_engine.get_running_strategy(strategy_id).paused is False
  141. finally:
  142. strategy_store.DB_PATH = original_db
  143. strategy_registry.STRATEGIES_DIR = original_dir
  144. strategy_engine._running.clear()