Download tests/test_fix_live.py from bryceweiner/fly-trader: direct link, hf CLI and curl.
- Browser
- Download file 21 kB
-
https://huggingface.co/bryceweiner/fly-trader/resolve/main/tests/test_fix_live.py
- Command line
-
hf download hf://bryceweiner/fly-trader/tests/test_fix_live.py
-
curl -L -o test_fix_live.py https://huggingface.co/bryceweiner/fly-trader/resolve/main/tests/test_fix_live.py
21 kB
| """Review fixes: partial-sell realized P&L, training reset leaves the circuit armed, close guard + runner lock, | |
| confirmation polling past lastValidBlockHeight.""" | |
| import threading | |
| import uuid | |
| import pytest | |
| from fly_trader import config | |
| from fly_trader.agent import runner | |
| from fly_trader.db.connection import transaction | |
| from fly_trader.execution import broker_live, ledger | |
| from fly_trader.execution.broker_live import FillResult | |
| from fly_trader.execution.broker_paper import PaperBroker | |
| from fly_trader.execution.worker import ExecRequest, ExecutionWorker | |
| def _pos(conn, pid): | |
| return conn.execute("SELECT * FROM positions WHERE id=%s", (pid,)).fetchone() | |
| # ---- 1. partial sells accumulate realized P&L; paper cash stays consistent ---- | |
| def test_paper_partial_then_close_accumulates_realized(db_conn): | |
| book, mint = f"replay_t{uuid.uuid4().hex[:8]}", f"M{uuid.uuid4().hex[:10]}" | |
| pid = ledger.open_position(db_conn, book=book, mint=mint, pool=None, qty_raw=1_000_000, cost_sol=1.0, entry_price=1.0, | |
| decision_id=None, fees_sol=0.0) | |
| assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0) | |
| pb = PaperBroker(book) | |
| kw = dict(decision_id=None, res_quote_sol=100.0, mcap_sol=100.0, program_label=None, forced_kind=None) | |
| f1 = pb.sell(db_conn, position=ledger.open_positions(db_conn, book)[0], price=1.2, fraction=0.5, **kw) | |
| assert f1.ok and 0 < f1.sol_delta < 0.6 | |
| p = _pos(db_conn, pid) | |
| assert p["status"] == "open" and int(p["qty"]) == 500_000 and p["cost_sol"] == pytest.approx(0.5) | |
| assert p["realized_sol"] == pytest.approx(f1.sol_delta - 0.5) | |
| assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0 + f1.sol_delta) | |
| stale = ledger.open_positions(db_conn, book)[0] | |
| f2 = pb.sell(db_conn, position=stale, price=0.8, **kw) | |
| assert f2.ok and "realized" in f2.reason | |
| p = _pos(db_conn, pid) | |
| assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(f1.sol_delta + f2.sol_delta - 1.0) | |
| assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0 + f1.sol_delta + f2.sol_delta) | |
| # 4. the close guard: a second close of the same position books nothing and records no fill | |
| f3 = pb.sell(db_conn, position=stale, price=0.8, **kw) | |
| assert not f3.ok and f3.fill_id is None | |
| assert ledger.close_position(db_conn, position_id=pid, exit_price=1.0, proceeds_sol=5.0, fees_sol=0.0, | |
| decision_id=None, forced_kind=None) is None | |
| assert _pos(db_conn, pid)["realized_sol"] == pytest.approx(f1.sol_delta + f2.sol_delta - 1.0) | |
| n = db_conn.execute("SELECT count(*) AS n FROM fills WHERE book=%s AND mint=%s AND side='sell'", (book, mint)).fetchone()["n"] | |
| assert n == 2 | |
| class _FakeBroker: | |
| def __init__(self): | |
| self.next: FillResult | None = None | |
| self.kw: dict = {} | |
| def swap(self, conn, **kw): | |
| self.kw = kw | |
| return self.next | |
| def test_worker_partial_then_close_books_each_sale_once(): | |
| mint = f"M{uuid.uuid4().hex[:10]}" | |
| with transaction() as conn: | |
| circ = dict(conn.execute("SELECT fail_count, tripped FROM circuit_state WHERE id=1").fetchone()) | |
| pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=1.0, entry_price=1.0, | |
| decision_id=None, fees_sol=0.0) | |
| fb = _FakeBroker() | |
| w = ExecutionWorker(fb) | |
| try: | |
| req = ExecRequest(decision_id=None, mint=mint, pool=None, side="sell", amount_in=400_000, slippage_bps=100, | |
| max_slippage_bps=100, decimals=6, position_id=pid) | |
| fb.next = FillResult(True, None, None, "s1", "confirmed", 0, None, -400_000, 500_000_000, 1.25, 1) | |
| r1 = w._execute(req) | |
| assert r1.ok and r1.realized_sol == pytest.approx(0.1) | |
| fb.next = FillResult(True, None, None, "s2", "confirmed", 0, None, -600_000, 900_000_000, 1.5, 1) | |
| r2 = w._execute(req) | |
| assert r2.ok and r2.realized_sol == pytest.approx(0.3) | |
| fb.next = FillResult(True, None, None, "s3", "confirmed", 0, None, -1, 1, 1.0, 1) | |
| r3 = w._execute(req) | |
| assert r3.ok and r3.realized_sol is None | |
| with transaction() as conn: | |
| p = _pos(conn, pid) | |
| assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(0.4) and p["cost_sol"] == pytest.approx(0.6) | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE id=%s", (pid,)) | |
| conn.execute("UPDATE circuit_state SET fail_count=%s, tripped=%s WHERE id=1", (circ["fail_count"], circ["tripped"])) | |
| # ---- 2. a training reset never re-arms the circuit ---- | |
| def test_training_reset_leaves_circuit_alone(tmp_path, monkeypatch): | |
| from fly_trader.ops.reset import reset_training_state | |
| monkeypatch.setattr(config, "PG_ARCHIVE_DIR", tmp_path / "arch") | |
| monkeypatch.setattr(config, "BRAIN_DIR", tmp_path / "brain") | |
| cols = "fail_count, tripped, kill_switch, kill_reason, peak_wealth" | |
| with transaction() as conn: | |
| before = dict(conn.execute(f"SELECT {cols} FROM circuit_state WHERE id=1").fetchone()) | |
| conn.execute("UPDATE circuit_state SET fail_count=3, tripped=true, kill_switch=true, kill_reason='dd', peak_wealth=5.0 WHERE id=1") | |
| try: | |
| reset_training_state(reason="test", archive=False) | |
| with transaction() as conn: | |
| after = dict(conn.execute(f"SELECT {cols} FROM circuit_state WHERE id=1").fetchone()) | |
| assert after == {"fail_count": 3, "tripped": True, "kill_switch": True, "kill_reason": "dd", "peak_wealth": 5.0} | |
| finally: | |
| with transaction() as conn: | |
| conn.execute("UPDATE circuit_state SET fail_count=%(fail_count)s, tripped=%(tripped)s, kill_switch=%(kill_switch)s, " | |
| "kill_reason=%(kill_reason)s, peak_wealth=%(peak_wealth)s WHERE id=1", before) | |
| # ---- 4. one runner per database ---- | |
| def test_runner_lock_refuses_a_second_runner(monkeypatch): | |
| monkeypatch.setattr(runner, "setup", lambda name: None) | |
| ran = [] | |
| monkeypatch.setattr(runner, "_main", lambda stop_event=None: ran.append(1)) | |
| held = runner.acquire_runner_lock() | |
| try: | |
| with pytest.raises(RuntimeError, match="runner lock"): | |
| runner.acquire_runner_lock() | |
| with pytest.raises(RuntimeError, match="runner lock"): | |
| runner.main() | |
| assert ran == [] | |
| finally: | |
| held.close() | |
| runner.main() # lock free again; main releases it on return | |
| assert ran == [1] | |
| runner.acquire_runner_lock().close() | |
| # ---- 5. a processed signature is polled past lastValidBlockHeight ---- | |
| class _StatusRpc: | |
| def __init__(self, statuses): | |
| self.statuses = list(statuses) | |
| def get_signature_statuses(self, sigs): | |
| return [self.statuses.pop(0)] | |
| def get_block_height(self): | |
| return 10_000 | |
| def test_await_confirmation_keeps_polling_a_processed_signature(monkeypatch): | |
| monkeypatch.setattr(broker_live.time, "sleep", lambda s: None) | |
| proc = {"err": None, "confirmationStatus": "processed"} | |
| rpc = _StatusRpc([proc, proc, {"err": None, "confirmationStatus": "confirmed"}]) | |
| assert broker_live.await_confirmation(rpc, "sig", 100)[0] == "confirmed" | |
| rpc = _StatusRpc([proc, None]) | |
| assert broker_live.await_confirmation(rpc, "sig", 100)[0] == "expired" | |
| # ---- 6. live keeps up with paper: retryable failures, sweeps, a buy behind a sale, rent back, a clean stop ---- | |
| class _ScriptBroker: | |
| """Fills from a script, in order; ``gate`` holds each swap until set (an order in flight); every call recorded.""" | |
| def __init__(self, fills, gate=None): | |
| self.fills, self.gate = list(fills), gate | |
| self.calls, self.started = [], threading.Event() | |
| def swap(self, conn, **kw): | |
| self.started.set(); self.calls.append(kw) | |
| if self.gate is not None: | |
| self.gate.wait(10) | |
| return self.fills.pop(0) | |
| def circuit1(): | |
| with transaction() as conn: | |
| before = dict(conn.execute("SELECT fail_count, tripped, last_failure_ts FROM circuit_state WHERE id=1").fetchone()) | |
| conn.execute("UPDATE circuit_state SET fail_count = 0, tripped = false WHERE id = 1") | |
| yield | |
| with transaction() as conn: | |
| conn.execute("UPDATE circuit_state SET fail_count=%(fail_count)s, tripped=%(tripped)s, last_failure_ts=%(last_failure_ts)s WHERE id=1", before) | |
| def _req(mint, side, amount_in=1_000_000, **kw): | |
| return ExecRequest(decision_id=None, mint=mint, pool=None, side=side, amount_in=amount_in, slippage_bps=300, max_slippage_bps=500, decimals=6, **kw) | |
| def _fails(): | |
| with transaction() as conn: | |
| return conn.execute("SELECT fail_count FROM circuit_state WHERE id=1").fetchone()["fail_count"] | |
| def _fill(token_delta, lamports_delta, price=1e-7): | |
| return FillResult(True, None, None, "sig", "confirmed", 0, None, token_delta, lamports_delta, price, 1) | |
| def test_failed_buys_that_moved_nothing_may_go_again_and_a_price_that_ran_is_no_execution_failure(circuit1): | |
| mint = f"M{uuid.uuid4().hex[:10]}" | |
| fb = _FakeBroker(); w = ExecutionWorker(fb) | |
| try: | |
| fb.next = FillResult(False, 1, None, None, "price_ran", None, "price ran: quoted 1.1e-07 SOL per token, +10.0% over the decision's 1e-07", 0, 0, None, 1) | |
| r = w._execute(_req(mint, "buy", 10 ** 8, ref_price=1e-7)) | |
| assert not r.ok and r.retryable and r.error.startswith("price ran") and (fb.kw["ref_price"], fb.kw["decimals"]) == (1e-7, 6) and _fails() == 0 | |
| fb.next = FillResult(False, 2, None, None, "order_error", 1, "Insufficient funds", 0, 0, None, 1) | |
| r = w._execute(_req(mint, "buy", 10 ** 8)) | |
| assert not r.ok and r.retryable and _fails() == 1 # a failure, but nothing landed | |
| fb.next = FillResult(False, 3, 99, "sig", "confirmed", 0, None, 0, 0, None, 1) | |
| assert not w._execute(_req(mint, "buy", 10 ** 8)).retryable and _fails() == 2 # an unverified fill row: never sent again | |
| with transaction() as conn: | |
| ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.1, entry_price=1e-7, decision_id=None, fees_sol=0.0) | |
| fb.kw = {} | |
| r = w._execute(_req(mint, "buy", 10 ** 8)) # the sale queued before it left the position open | |
| assert not r.ok and r.retryable and r.error.startswith("held") and fb.kw == {} and _fails() == 2 | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_a_dead_bag_sweep_books_its_proceeds_onto_the_written_off_position(circuit1): | |
| mint = f"M{uuid.uuid4().hex[:10]}" | |
| with transaction() as conn: | |
| pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) | |
| ledger.close_position(conn, position_id=pid, exit_price=0.0, proceeds_sol=0.0, fees_sol=0.0, decision_id=None, forced_kind="dead_bag") | |
| fb = _FakeBroker(); w = ExecutionWorker(fb) | |
| try: | |
| fb.next = _fill(-1_000_000, 30_000_000, 3e-8) | |
| r = w._execute(_req(mint, "sell", forced_kind="dead_bag_sweep")) | |
| assert r.ok and r.position_id == pid and r.closed and r.realized_sol == pytest.approx(0.03) | |
| with transaction() as conn: | |
| p = _pos(conn, pid) | |
| assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(-0.2 + 0.03) | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_a_sale_that_empties_its_mint_leaves_the_token_accounts_close_to_the_mirror(circuit1): | |
| """Live money first: the worker goes on to the next order; the mirror closes the emptied account mid-minute.""" | |
| mint = f"M{uuid.uuid4().hex[:10]}" | |
| with transaction() as conn: | |
| pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) | |
| fb = _ScriptBroker([_fill(-1_000_000, 250_000_000, 2.5e-7)]); w = ExecutionWorker(fb) | |
| try: | |
| assert w.submit(_req(mint, "sell", position_id=pid)) | |
| r = w.results.get(timeout=10) | |
| assert r.ok and r.closed and w.emptied == {mint} and r.realized_sol == pytest.approx(0.05) | |
| with transaction() as conn: | |
| p = _pos(conn, pid) | |
| assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(0.25 - 0.2) and w.pending() == set() | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_a_first_buys_token_account_rent_is_a_deposit_its_close_returns(circuit1): | |
| """The rent a first buy pays for the mint's token account is kept apart from the position's cost (a live buy costs | |
| what its paper twin does) and counted in live wealth while held; the sale's close returns it, and only the close's | |
| fee is the position's.""" | |
| mint = f"M{uuid.uuid4().hex[:10]}"; rent, fee = 1_513_840, 5_000 | |
| buy = FillResult(True, None, None, "sig", "confirmed", 0, None, 1_000_000, -200_000_000, 2e-7, 1, rent) | |
| fb = _ScriptBroker([buy, _fill(-1_000_000, 250_000_000, 2.5e-7)]); w = ExecutionWorker(fb) | |
| try: | |
| assert w.submit(_req(mint, "buy", 200_000_000)) | |
| r1 = w.results.get(timeout=10) | |
| with transaction() as conn: | |
| p = _pos(conn, r1.position_id); held = ledger.deposits_held(conn, "live") | |
| assert r1.ok and p["cost_sol"] == pytest.approx(0.2) and p["deposit_sol"] == pytest.approx(rent / 1e9) and held >= rent / 1e9 | |
| assert w.submit(_req(mint, "sell", position_id=r1.position_id)) | |
| r2 = w.results.get(timeout=10) | |
| with transaction() as conn: | |
| before = ledger.deposits_held(conn, "live") | |
| booked = ledger.book_account_close(conn, "live", mint, rent - fee) # the mirror's close (agent/fly_live.close_accounts) | |
| p = _pos(conn, r1.position_id); after = ledger.deposits_held(conn, "live") | |
| assert r2.closed and w.emptied == {mint} and r2.realized_sol == pytest.approx(0.05) | |
| assert booked == (pytest.approx(-fee / 1e9), r1.position_id) and before - after == pytest.approx(rent / 1e9) | |
| assert p["deposit_refunded"] and p["realized_sol"] == pytest.approx(0.25 - 0.2 - fee / 1e9) | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_a_multi_hop_buy_has_the_mirror_close_its_intermediate_token_accounts(circuit1): | |
| """The accounts a multi-hop route opened for its intermediate tokens and left empty are the mirror's to close | |
| mid-minute, as a sale's are: their rent comes back. It is neither the position's cost nor its deposit.""" | |
| mint, hop = f"M{uuid.uuid4().hex[:10]}", f"H{uuid.uuid4().hex[:10]}" | |
| buy = FillResult(True, None, None, "sig", "confirmed", 0, None, 1_000_000, -200_000_000, 2e-7, 1, 0, {hop: 1_488_440}) | |
| fb = _ScriptBroker([buy]); w = ExecutionWorker(fb) | |
| try: | |
| assert w.submit(_req(mint, "buy", 200_000_000)) | |
| r = w.results.get(timeout=10) | |
| with transaction() as conn: | |
| p = _pos(conn, r.position_id) | |
| assert r.ok and w.emptied == {hop} and p["cost_sol"] == pytest.approx(0.2) and p["deposit_sol"] is None | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_a_buy_queues_behind_the_sale_of_its_mint_which_keeps_its_token_account(circuit1): | |
| """Paper sold the token and bought it back in one minute: the live buy waits behind the sale, runs once the sale has | |
| closed the position, and the token account stays for it.""" | |
| mint = f"M{uuid.uuid4().hex[:10]}" | |
| with transaction() as conn: | |
| pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) | |
| gate = threading.Event() | |
| fb = _ScriptBroker([_fill(-1_000_000, 250_000_000, 2.5e-7), _fill(900_000, -200_000_000, 2.2e-7)], gate=gate); w = ExecutionWorker(fb) | |
| try: | |
| assert w.submit(_req(mint, "sell", position_id=pid)) and fb.started.wait(10) and w.selling() == {mint} | |
| assert w.submit(_req(mint, "buy", 200_000_000)) | |
| assert not w.submit(_req(mint, "buy", 1)) and not w.submit(_req(mint, "sell")) # one buy behind one sale, nothing more | |
| assert w.selling() == set() and w.pending() == {mint} | |
| gate.set() | |
| r1, r2 = w.results.get(timeout=10), w.results.get(timeout=10) | |
| assert (r1.request.side, r1.ok, r1.closed, r2.request.side, r2.ok) == ("sell", True, True, "buy", True) and w.emptied == set() | |
| with transaction() as conn: | |
| rows = conn.execute("SELECT status, qty FROM positions WHERE book = 'live' AND mint = %s ORDER BY id", (mint,)).fetchall() | |
| assert [(r["status"], int(r["qty"])) for r in rows] == [("closed", 1_000_000), ("open", 900_000)] | |
| finally: | |
| gate.set(); w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |
| def test_stop_waits_for_the_order_in_flight_and_drops_the_queued_ones(circuit1): | |
| a, b = f"M{uuid.uuid4().hex[:10]}", f"M{uuid.uuid4().hex[:10]}" | |
| gate = threading.Event() | |
| fb = _ScriptBroker([_fill(-5, 5), _fill(-5, 5)], gate=gate); w = ExecutionWorker(fb, threads=1) # one thread: b waits in the queue | |
| assert w.submit(_req(a, "sell")) and fb.started.wait(10) and w.submit(_req(b, "sell")) | |
| threading.Timer(0.3, gate.set).start() | |
| w.stop(timeout=10) | |
| assert not w.thread.is_alive() and [c["mint"] for c in fb.calls] == [a] and w.pending() == set() | |
| assert w.results.get_nowait().request.mint == a and not w.submit(_req(b, "sell")) | |
| def test_swaps_of_different_mints_run_at_once_and_each_mints_in_order(circuit1): | |
| """Live money fast: one token's swap in flight never holds another token's back; one token's sale and buy-back still | |
| run one after the other.""" | |
| a, b = f"M{uuid.uuid4().hex[:10]}", f"M{uuid.uuid4().hex[:10]}" | |
| gate = threading.Event() | |
| fb = _ScriptBroker([_fill(-5, 5), _fill(-5, 5), _fill(5, -5)], gate=gate); w = ExecutionWorker(fb, threads=3) | |
| try: | |
| assert w.submit(_req(a, "sell")) and w.submit(_req(b, "sell")) and w.submit(_req(a, "buy", 5)) | |
| for _ in range(100): | |
| if len(fb.calls) >= 2: | |
| break | |
| threading.Event().wait(0.02) | |
| assert sorted((c["mint"], c["side"]) for c in fb.calls) == sorted([(a, "sell"), (b, "sell")]) # both sales in flight, a's buy behind its sale | |
| gate.set() | |
| got = [w.results.get(timeout=10) for _ in range(3)] | |
| assert [(c["mint"], c["side"]) for c in fb.calls if c["mint"] == a] == [(a, "sell"), (a, "buy")] # a's buy ran once its sale was done | |
| assert all(r.ok for r in got) and w.pending() == set() | |
| finally: | |
| gate.set(); w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = ANY(%s)", ([a, b],)) | |
| def test_a_wealth_mark_waits_for_every_swap_in_flight_and_never_holds_one_back(): | |
| from fly_trader.execution.worker import SharedLock | |
| busy = SharedLock() | |
| with busy.shared(): | |
| with busy.shared(): # two swaps in flight at once | |
| assert not busy.acquire(blocking=False) and not busy.acquire(timeout=0.05) | |
| assert busy.acquire(timeout=0.1) # none in flight: the mark reads alone | |
| started = threading.Event() | |
| def swap(): | |
| with busy.shared(): | |
| started.set() | |
| t = threading.Thread(target=swap); t.start() | |
| assert not started.wait(0.1) # a swap waits only while a mark holds it | |
| busy.release(); t.join(2) | |
| assert started.is_set() | |
| def test_the_worker_tells_its_mirror_once_an_order_is_booked(circuit1): | |
| """The live stats follow the live orders: the mirror's hook runs as soon as each order is booked, after its result.""" | |
| mint = f"M{uuid.uuid4().hex[:10]}"; seen = [] | |
| fb = _ScriptBroker([_fill(1_000_000, -200_000_000, 2e-7)]); w = ExecutionWorker(fb) | |
| w.on_done = lambda res: seen.append((res.request.side, res.ok, res.position_id is not None)) | |
| try: | |
| assert w.submit(_req(mint, "buy", 200_000_000)) | |
| w.results.get(timeout=10) | |
| for _ in range(100): | |
| if seen: | |
| break | |
| threading.Event().wait(0.05) | |
| assert seen == [("buy", True, True)] | |
| finally: | |
| w.stop() | |
| with transaction() as conn: | |
| conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) | |