ai-exchange/tests/unit/test_runtime_script_json.py

import contextlib
import io
import json
import os
import unittest
from unittest.mock import patch

from sqlalchemy.exc import OperationalError

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

from scripts import (
    apply_migrations,
    check_db,
    check_db_connection,
    run_api_server,
    run_dev,
    run_displacement,
    run_fvg,
    run_liquidity,
    run_mss,
    run_ranges,
    run_session_generator,
    run_sweeps,
    run_swings,
    run_training_cases,
    show_models,
)


class _FakeTrainingCaseGenerator:
    def __init__(self) -> None:
        self.calls = []

    def generate_cases(self, instrument_id, timeframe, case_types=None, limit_per_type=10):
        self.calls.append(
            {
                "instrument_id": instrument_id,
                "timeframe": timeframe,
                "case_types": list(case_types),
                "limit_per_type": limit_per_type,
            }
        )
        return run_training_cases.CaseGenerationResult(
            generated=2,
            by_type={"sweep": 1, "range": 1},
            case_ids=[101, 202],
        )


class RuntimeScriptJsonTests(unittest.TestCase):
    def test_structure_scripts_report_database_unavailable_json(self) -> None:
        cases = [
            (run_session_generator, "SESSION_GENERATOR_OUTPUT_JSON", "scripts.run_session_generator.SessionGenerator.generate_for_date"),
            (run_swings, "SWINGS_OUTPUT_JSON", "scripts.run_swings.SwingService.detect_swings"),
            (run_ranges, "RANGES_OUTPUT_JSON", "scripts.run_ranges.DealingRangeService.build_active_range"),
            (run_liquidity, "LIQUIDITY_OUTPUT_JSON", "scripts.run_liquidity.LiquidityPoolService.build_candidates"),
            (run_sweeps, "SWEEPS_OUTPUT_JSON", "scripts.run_sweeps.SweepService.detect_sweeps"),
            (run_displacement, "DISPLACEMENT_OUTPUT_JSON", "scripts.run_displacement.DisplacementService.detect_displacement"),
            (run_mss, "MSS_OUTPUT_JSON", "scripts.run_mss.MSSService.detect_mss"),
            (run_fvg, "FVG_OUTPUT_JSON", "scripts.run_fvg.FVGService.detect_fvg"),
        ]

        for module, env_key, patch_target in cases:
            with self.subTest(script=module.__name__):
                output = io.StringIO()
                with patch(patch_target, side_effect=OperationalError("stmt", {}, Exception())), contextlib.redirect_stdout(output):
                    exit_code = module.main(env={env_key: "1"})

                payload = json.loads(output.getvalue())
                self.assertEqual(exit_code, 2)
                self.assertFalse(payload["passed"])
                self.assertEqual(payload["error_type"], "OperationalError")
                self.assertEqual(payload["instrument_id"], 1)
                self.assertEqual(payload["timeframe"], "1m")

    def test_database_check_scripts_report_unavailable_json(self) -> None:
        cases = [
            (check_db, "CHECK_DB_OUTPUT_JSON", "scripts.check_db.engine.connect"),
            (check_db_connection, "CHECK_DB_CONNECTION_OUTPUT_JSON", "scripts.check_db_connection.engine.connect"),
        ]

        for module, env_key, patch_target in cases:
            with self.subTest(script=module.__name__):
                output = io.StringIO()
                with patch(patch_target, side_effect=OperationalError("stmt", {}, Exception())), contextlib.redirect_stdout(output):
                    exit_code = module.main(env={env_key: "1"})

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

    def test_training_cases_report_database_unavailable_json(self) -> None:
        output = io.StringIO()
        with patch(
            "scripts.run_training_cases.TrainingCaseGenerator.generate_cases",
            side_effect=OperationalError("stmt", {}, Exception()),
        ), contextlib.redirect_stdout(output):
            exit_code = run_training_cases.main(env={"TRAINING_CASES_OUTPUT_JSON": "1"})

        payload = json.loads(output.getvalue())
        self.assertEqual(exit_code, 2)
        self.assertFalse(payload["passed"])
        self.assertEqual(payload["error_type"], "OperationalError")
        self.assertEqual(payload["instrument_id"], 1)
        self.assertEqual(payload["timeframe"], "1m")
        self.assertEqual(payload["case_types"], ["sweep", "fvg", "range", "mss_displacement"])
        self.assertEqual(payload["generated"], 0)

    def test_training_cases_outputs_generation_json(self) -> None:
        output = io.StringIO()
        generator = _FakeTrainingCaseGenerator()

        with contextlib.redirect_stdout(output):
            exit_code = run_training_cases.run_training_generation(
                generator=generator,
                instrument_id=7,
                timeframe="5m",
                case_types=["sweep", "range"],
                limit_per_type=3,
                output_json=True,
            )

        payload = json.loads(output.getvalue())
        self.assertEqual(exit_code, 0)
        self.assertTrue(payload["passed"])
        self.assertEqual(payload["instrument_id"], 7)
        self.assertEqual(payload["timeframe"], "5m")
        self.assertEqual(payload["limit_per_type"], 3)
        self.assertEqual(payload["case_types"], ["sweep", "range"])
        self.assertEqual(payload["generated"], 2)
        self.assertEqual(payload["by_type"], {"range": 1, "sweep": 1})
        self.assertEqual(payload["case_ids"], [101, 202])
        self.assertEqual(generator.calls[0]["case_types"], ["sweep", "range"])

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

        with contextlib.redirect_stdout(output):
            exit_code = show_models.main(env={"SHOW_MODELS_OUTPUT_JSON": "1"})

        payload = json.loads(output.getvalue())
        self.assertEqual(exit_code, 0)
        self.assertTrue(payload["passed"])
        self.assertGreater(payload["model_count"], 0)
        self.assertIn("candle", payload["models"])

    def test_apply_migrations_delegates_to_safe_runner(self) -> None:
        with patch("scripts.apply_migrations.run_migrations_main", return_value=2) as run_migrations:
            exit_code = apply_migrations.main(env={"MIGRATIONS_OUTPUT_JSON": "1"})

        self.assertEqual(exit_code, 2)
        run_migrations.assert_called_once_with(env={"MIGRATIONS_OUTPUT_JSON": "1"})

    def test_run_dev_preserves_zero_exit_code(self) -> None:
        with patch("scripts.run_dev.app_main") as app_main:
            exit_code = run_dev.main()

        self.assertEqual(exit_code, 0)
        app_main.assert_called_once_with()

    def test_run_api_server_uses_env_and_returns_when_server_stops(self) -> None:
        with patch("scripts.run_api_server.serve_api") as serve_api:
            exit_code = run_api_server.main(env={"AI_ICT_API_HOST": "127.0.0.1", "AI_ICT_API_PORT": "8123"})

        self.assertEqual(exit_code, 0)
        serve_api.assert_called_once_with(host="127.0.0.1", port=8123)


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