xiaobaifupan/app/tests/test_preservation_slice_screener.py

from __future__ import annotations

import ast
import hashlib
import unittest
from pathlib import Path

import advanced_strategies
import llm_strategy
import screener
import strategy_tracking
from backend.features.screener import compiler, engine, strategies, tracking
from backend.features.screener import service as screener_service
from tests.preservation_helpers import (
    assert_frontend_runtime_matches_audited_baseline,
    assert_moved_asset_matches,
    assert_page_prefix_matches,
)


APP_ROOT = Path(__file__).resolve().parents[1]
ORIGINAL_ROOT = APP_ROOT.parent

SCREENER_SERVICE_METHODS = {
    "_strategy_missing_data",
    "screener_setup",
    "screener_tracking",
    "add_screener_tracking",
    "remove_screener_tracking",
    "refresh_screener_tracking",
    "sync_screener_data",
    "_schedule_automatic_screeners",
    "run_automatic_screeners",
    "compile_screener_strategy",
    "save_screener_strategy",
    "delete_screener_strategy",
    "run_screener",
}

SCREENER_REPOSITORY_METHODS = {
    "upsert_benchmark_bars",
    "upsert_daily_indicators",
    "upsert_fundamental_indicators",
    "upsert_moneyflow",
    "upsert_earnings_events",
    "daily_indicator_dates",
    "fundamental_periods",
    "factor_dates",
    "factor_health_summary",
    "load_factor_data",
    "snapshot_summaries",
    "save_screener_strategy",
    "list_screener_strategies",
    "delete_screener_strategy",
    "save_screener_run",
    "_screener_run_payload",
    "latest_screener_run",
    "latest_screener_runs",
    "latest_screener_context_runs",
    "screener_runs_for_date",
    "get_screener_run",
    "save_strategy_tracks",
    "list_strategy_tracks",
    "delete_strategy_track",
    "load_tracking_bars",
}


def class_methods(path: Path, class_name: str) -> dict[str, str]:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    owner = next(
        node
        for node in tree.body
        if isinstance(node, ast.ClassDef) and node.name == class_name
    )
    return {
        node.name: ast.dump(node, include_attributes=False)
        for node in owner.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
    }


def top_level_definition(path: Path, name: str) -> str:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    node = next(
        item
        for item in tree.body
        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
        and item.name == name
    )
    return ast.dump(node, include_attributes=False)


def module_without_imports(path: Path) -> str:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    tree.body = [
        node for node in tree.body if not isinstance(node, (ast.Import, ast.ImportFrom))
    ]
    return ast.dump(tree, include_attributes=False)


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


class ScreenerSliceSourceEquivalenceTests(unittest.TestCase):
    def assert_methods_equal(
        self,
        original_path: Path,
        original_class: str,
        migrated_path: Path,
        migrated_class: str,
        names: set[str],
    ) -> None:
        original = class_methods(original_path, original_class)
        migrated = class_methods(migrated_path, migrated_class)
        self.assertEqual(set(migrated), names)
        for name in sorted(names):
            self.assertEqual(migrated[name], original[name], name)

    def test_screener_service_is_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "server.py",
            "DashboardService",
            APP_ROOT / "backend" / "features" / "screener" / "service.py",
            "ScreenerServiceMixin",
            SCREENER_SERVICE_METHODS,
        )
        self.assertEqual(
            top_level_definition(ORIGINAL_ROOT / "server.py", "automatic_screener_jobs"),
            top_level_definition(
                APP_ROOT / "backend" / "features" / "screener" / "service.py",
                "automatic_screener_jobs",
            ),
        )
        self.assertEqual(screener_service.SCREENER_LIBRARY_VERSION, 8)

    def test_screener_repository_is_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "database.py",
            "ReviewDatabase",
            APP_ROOT / "backend" / "features" / "screener" / "repository.py",
            "ScreenerRepositoryMixin",
            SCREENER_REPOSITORY_METHODS,
        )

    def test_moved_methods_are_not_duplicated(self) -> None:
        service_methods = class_methods(
            APP_ROOT / "backend" / "application.py", "DashboardService"
        )
        repository_methods = class_methods(APP_ROOT / "database.py", "ReviewDatabase")
        self.assertTrue(SCREENER_SERVICE_METHODS.isdisjoint(service_methods))
        self.assertTrue(SCREENER_REPOSITORY_METHODS.isdisjoint(repository_methods))

    def test_engine_and_tracking_logic_match_the_original(self) -> None:
        self.assertEqual(
            module_without_imports(ORIGINAL_ROOT / "screener.py"),
            module_without_imports(
                APP_ROOT / "backend" / "features" / "screener" / "engine.py"
            ),
        )
        self.assertEqual(
            class_methods(
                ORIGINAL_ROOT / "backend" / "features" / "screener" / "tracking.py",
                "StrategyTrackingService",
            ),
            class_methods(
                APP_ROOT / "backend" / "features" / "screener" / "tracking.py",
                "StrategyTrackingService",
            ),
        )

    def test_library_is_exact_and_compiler_uses_shared_transport(self) -> None:
        self.assertEqual(
            sha256(ORIGINAL_ROOT / "advanced_strategies.py"),
            sha256(APP_ROOT / "backend/features/screener/strategies.py"),
        )
        compiler_source = (
            APP_ROOT / "backend/features/screener/compiler.py"
        ).read_text(encoding="utf-8")
        self.assertEqual(compiler_source.count("llm_transport.chat_completion"), 2)
        self.assertNotIn("urllib.request", compiler_source)

    def test_compatibility_modules_export_the_canonical_objects(self) -> None:
        self.assertIs(screener, engine)
        self.assertIs(advanced_strategies, strategies)
        self.assertIs(llm_strategy, compiler)
        self.assertIs(
            strategy_tracking.StrategyTrackingService,
            tracking.StrategyTrackingService,
        )

    def test_screener_frontend_assets_are_unchanged(self) -> None:
        assert_frontend_runtime_matches_audited_baseline(self)
        assert_moved_asset_matches(self, "styles.css", "styles/styles.css")
        assert_page_prefix_matches(self, "pages/screener/page.js")


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