ai-exchange/tests/unit/test_okx_backfill_script.py

import contextlib
import io
import json
import os
import unittest
from urllib.error import URLError
from unittest.mock import patch

os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://user:pass@localhost:5432/ai_ict_test")

from scripts import run_okx_backfill
from src.services.market_data import OKXCandle


class OKXBackfillScriptTests(unittest.TestCase):
    def test_build_backfill_config_from_env_uses_defaults(self) -> None:
        config = run_okx_backfill.build_backfill_config_from_env({})

        self.assertEqual(config["symbol"], "BTC-USDT")
        self.assertEqual(config["inst_id"], "BTC-USDT-SWAP")
        self.assertEqual(config["pages"], 20)
        self.assertEqual(config["limit"], 100)
        self.assertEqual(config["timeframes"], ["1m", "5m", "15m"])

    def test_build_backfill_config_from_env_reads_overrides(self) -> None:
        config = run_okx_backfill.build_backfill_config_from_env(
            {
                "OKX_BACKFILL_SYMBOL": "ETH-USDT",
                "OKX_BACKFILL_INST_ID": "ETH-USDT-SWAP",
                "OKX_BACKFILL_PAGES": "3",
                "OKX_BACKFILL_LIMIT": "50",
                "OKX_BACKFILL_TIMEFRAMES": "1m,15m",
            }
        )

        self.assertEqual(config["symbol"], "ETH-USDT")
        self.assertEqual(config["inst_id"], "ETH-USDT-SWAP")
        self.assertEqual(config["pages"], 3)
        self.assertEqual(config["limit"], 50)
        self.assertEqual(config["timeframes"], ["1m", "15m"])

    def test_run_okx_backfill_once_deduplicates_confirmed_candles(self) -> None:
        client = FakeOKXClient(
            [
                OKXCandle("BTC-USDT-SWAP", "1m", "1000", "1", "2", "0.5", "1.5", "10", "1"),
                OKXCandle("BTC-USDT-SWAP", "1m", "1000", "1", "2", "0.5", "1.5", "10", "1"),
                OKXCandle("BTC-USDT-SWAP", "1m", "2000", "2", "3", "1.5", "2.5", "11", "0"),
            ]
        )
        repository = FakeCandleRepository()

        summary = run_okx_backfill.run_okx_backfill_once(
            client=client,
            repository=repository,
            pages=1,
            limit=3,
            timeframes=["1m"],
        )

        self.assertEqual(summary["total_fetched"], 3)
        self.assertEqual(summary["total_wrote"], 1)
        self.assertEqual(repository.batch_sizes, [1])

    def test_run_okx_backfill_once_uses_after_cursor_for_older_pages(self) -> None:
        client = CursorOKXClient()
        repository = FakeCandleRepository()

        summary = run_okx_backfill.run_okx_backfill_once(
            client=client,
            repository=repository,
            pages=2,
            limit=2,
            timeframes=["1m"],
        )

        self.assertEqual(client.after_values, [None, "2000"])
        self.assertEqual(summary["total_fetched"], 3)
        self.assertEqual(summary["total_wrote"], 3)

    def test_main_returns_failure_when_backfill_fails(self) -> None:
        output = io.StringIO()

        with patch("scripts.run_okx_backfill.run_okx_backfill_once", side_effect=URLError("offline")), contextlib.redirect_stdout(output):
            exit_code = run_okx_backfill.main()

        self.assertEqual(exit_code, 2)
        self.assertIn("OKX backfill failed: URLError", output.getvalue())

    def test_main_can_emit_json_summary(self) -> None:
        output = io.StringIO()
        summary = {
            "symbol": "BTC-USDT",
            "inst_id": "BTC-USDT-SWAP",
            "pages": 1,
            "limit": 2,
            "timeframes": [{"timeframe": "1m", "fetched": 2, "wrote": 2}],
            "total_fetched": 2,
            "total_wrote": 2,
        }

        with patch("scripts.run_okx_backfill.run_okx_backfill_once", return_value=summary), contextlib.redirect_stdout(output):
            exit_code = run_okx_backfill.main(env={"OKX_BACKFILL_OUTPUT_JSON": "1"})

        payload = json.loads(output.getvalue())
        self.assertEqual(exit_code, 0)
        self.assertTrue(payload["passed"])
        self.assertEqual(payload["total_wrote"], 2)

    def test_main_can_emit_json_failure(self) -> None:
        output = io.StringIO()

        with patch("scripts.run_okx_backfill.run_okx_backfill_once", side_effect=ValueError("bad")), contextlib.redirect_stdout(output):
            exit_code = run_okx_backfill.main(env={"OKX_BACKFILL_OUTPUT_JSON": "1"})

        payload = json.loads(output.getvalue())
        self.assertEqual(exit_code, 2)
        self.assertFalse(payload["passed"])
        self.assertEqual(payload["error_type"], "ValueError")


class FakeOKXClient:
    def __init__(self, candles: list[OKXCandle]) -> None:
        self.candles = candles

    def get_candles(self, **kwargs) -> list[OKXCandle]:
        return self.candles


class CursorOKXClient:
    def __init__(self) -> None:
        self.after_values = []

    def get_candles(self, **kwargs) -> list[OKXCandle]:
        after = kwargs.get("after")
        self.after_values.append(after)
        if after is None:
            return [
                OKXCandle("BTC-USDT-SWAP", "1m", "3000", "1", "2", "0.5", "1.5", "10", "1"),
                OKXCandle("BTC-USDT-SWAP", "1m", "2000", "1", "2", "0.5", "1.5", "10", "1"),
            ]
        return [
            OKXCandle("BTC-USDT-SWAP", "1m", "1000", "1", "2", "0.5", "1.5", "10", "1"),
        ]


class FakeCandleRepository:
    def __init__(self) -> None:
        self.batch_sizes = []

    def upsert_many(self, candles, instrument_id: int) -> int:
        self.batch_sizes.append(len(candles))
        return len(candles)


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