124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
from datetime import time
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import pytest
|
|
|
|
from llm_trade_bot.agent import AgentTools, _has_existing_idea
|
|
from llm_trade_bot.config import (
|
|
AppConfig,
|
|
BraveConfig,
|
|
BybitConfig,
|
|
ZaiConfig,
|
|
PromptsConfig,
|
|
RiskConfig,
|
|
ScannerConfig,
|
|
ScheduleConfig,
|
|
StorageConfig,
|
|
TelegramConfig,
|
|
)
|
|
from llm_trade_bot.database import Database
|
|
from llm_trade_bot.models import AccountSnapshot, InstrumentRules
|
|
|
|
|
|
class FakeBybit:
|
|
def __init__(self):
|
|
self.orders = []
|
|
|
|
async def account_state(self):
|
|
return {"wallet": [{"coin": []}], "positions": [], "open_orders": []}
|
|
|
|
async def market_context(self, category, symbol):
|
|
return {"ticker": {"lastPrice": "100"}}
|
|
|
|
async def account_snapshot(self):
|
|
return AccountSnapshot(Decimal(10000), Decimal(10000), Decimal(0), 0)
|
|
|
|
async def instrument_rules(self, category, symbol):
|
|
return InstrumentRules(Decimal("0.1"), Decimal("0.001"), Decimal("0.001"), Decimal(5), Decimal(100))
|
|
|
|
async def set_leverage(self, symbol, leverage):
|
|
return {"retCode": 0}
|
|
|
|
async def place_order(self, **parameters):
|
|
self.orders.append(parameters)
|
|
return {"retCode": 0, "result": {"orderId": f"order-{len(self.orders)}"}}
|
|
|
|
|
|
def make_config(database_path: Path) -> AppConfig:
|
|
return AppConfig(
|
|
ZaiConfig("key"),
|
|
BybitConfig("key", "secret"),
|
|
TelegramConfig(1, "hash", "token", frozenset({1})),
|
|
BraveConfig("key"),
|
|
PromptsConfig(database_path.parent / "prompts.md"),
|
|
ScheduleConfig(ZoneInfo("Europe/Warsaw"), (time(9), time(21))),
|
|
ScannerConfig(),
|
|
RiskConfig(),
|
|
StorageConfig(database_path),
|
|
)
|
|
|
|
|
|
async def test_open_trade_is_risk_sized_and_deduplicated_within_run(tmp_path):
|
|
database = Database(tmp_path / "state.db")
|
|
await database.open()
|
|
bybit = FakeBybit()
|
|
tools = AgentTools(make_config(database.path), database, bybit) # type: ignore[arg-type]
|
|
run_id = await database.start_run("test")
|
|
context = {"run_id": run_id, "candidate_keys": {"linear:BTCUSDT"}}
|
|
try:
|
|
first = await tools.open_trade(
|
|
"linear", "BTCUSDT", "Buy", "Market", 90, 2, "test", take_profit=120, leverage=3, context=context
|
|
)
|
|
second = await tools.open_trade(
|
|
"linear", "BTCUSDT", "Buy", "Market", 90, 2, "test", take_profit=120, leverage=3, context=context
|
|
)
|
|
assert first["quantity"] == "20"
|
|
assert second["deduplicated"] is True
|
|
assert len(bybit.orders) == 1
|
|
assert bybit.orders[0]["stopLoss"] == "90"
|
|
finally:
|
|
await database.close()
|
|
|
|
|
|
async def test_pause_blocks_entry_from_running_session(tmp_path):
|
|
database = Database(tmp_path / "state.db")
|
|
await database.open()
|
|
tools = AgentTools(make_config(database.path), database, FakeBybit()) # type: ignore[arg-type]
|
|
await database.set_state("paused", "true")
|
|
try:
|
|
with pytest.raises(ValueError, match="paused"):
|
|
await tools.open_trade(
|
|
"linear",
|
|
"BTCUSDT",
|
|
"Buy",
|
|
"Market",
|
|
90,
|
|
2,
|
|
"test",
|
|
take_profit=120,
|
|
context={"run_id": 1, "candidate_keys": {"linear:BTCUSDT"}},
|
|
)
|
|
finally:
|
|
await database.close()
|
|
|
|
|
|
def test_seeded_spot_balance_does_not_block_new_agent_entry():
|
|
state = {
|
|
"wallets": {"spot": {"balances": [{"asset": "SOL", "free": "4", "locked": "0"}]}},
|
|
"positions": [],
|
|
"open_orders": [],
|
|
}
|
|
assert _has_existing_idea(state, "spot", "SOLUSDT") is False
|
|
|
|
state["open_orders"] = [
|
|
{
|
|
"category": "spot",
|
|
"symbol": "SOLUSDT",
|
|
"protective": True,
|
|
"client_id": "ltb-12-abcd-tp",
|
|
}
|
|
]
|
|
assert _has_existing_idea(state, "spot", "SOLUSDT") is True
|