xiaobaifupan/app/tests/test_preservation_slice_heaven.py

from __future__ import annotations

import ast
import hashlib
import unittest
from pathlib import Path

import heaven_agent
import heaven_engine
from backend.features.heaven import agent as canonical_agent
from backend.features.heaven import engine as canonical_engine


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

HEAVEN_SERVICE_METHODS = {
    "_heaven_manual_schema",
    "_validate_heaven_manual_data",
    "_apply_heaven_manual_data",
    "_heaven_line_checks",
    "_resolve_heaven_stock_code",
    "heaven_setup",
    "_heaven_stock_context",
    "_heaven_market_mode",
    "_heaven_trend_sources",
    "_heaven_trend_quality_issues",
    "heaven_personal",
    "heaven_hexagram",
    "heaven_readings",
    "_heaven_reading_identity",
    "heaven_interpret",
    "_legacy_truncated_heaven_reading",
    "_call_heaven_agent",
    "_heaven_index_context",
    "_aggregate_index_context",
    "_heaven_sector_context",
}

HEAVEN_REPOSITORY_METHODS = {
    "_heaven_reading_dict",
    "save_heaven_reading",
    "list_heaven_readings",
    "latest_heaven_reading",
    "delete_heaven_reading",
}

HEAVEN_HTTP_METHODS = {"heaven_hexagram", "heaven_personal", "heaven_interpret"}


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


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_definitions(path: Path) -> dict[str, str]:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    return {
        node.name: ast.dump(node, include_attributes=False)
        for node in tree.body
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
    }


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

    def test_heaven_agent_uses_shared_transport(self) -> None:
        source = (
            APP_ROOT / "backend" / "features" / "heaven" / "agent.py"
        ).read_text(encoding="utf-8")
        self.assertIn("llm_transport.chat_completion", source)
        self.assertNotIn("urllib.request", source)

    def test_heaven_engine_definitions_are_exact_original_ast(self) -> None:
        self.assertEqual(
            top_level_definitions(ORIGINAL_ROOT / "heaven_engine.py"),
            top_level_definitions(
                APP_ROOT / "backend" / "features" / "heaven" / "engine.py"
            ),
        )

    def test_compatibility_modules_are_canonical_module_objects(self) -> None:
        self.assertIs(heaven_agent, canonical_agent)
        self.assertIs(heaven_engine, canonical_engine)

    def test_heaven_service_methods_are_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "server.py",
            "DashboardService",
            APP_ROOT / "backend" / "features" / "heaven" / "service.py",
            "HeavenServiceMixin",
            HEAVEN_SERVICE_METHODS,
            {"_heaven_reading_identity"},
        )
        source = (
            APP_ROOT / "backend" / "features" / "heaven" / "service.py"
        ).read_text(encoding="utf-8")
        self.assertIn(
            "MarketServiceMixin._display_compact_date(context_date)", source
        )

    def test_heaven_repository_methods_are_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "database.py",
            "ReviewDatabase",
            APP_ROOT / "backend" / "features" / "heaven" / "repository.py",
            "HeavenRepositoryMixin",
            HEAVEN_REPOSITORY_METHODS,
        )

    def test_original_classes_no_longer_duplicate_moved_methods(self) -> None:
        remaining_service = class_methods(
            APP_ROOT / "backend" / "application.py", "DashboardService"
        )
        remaining_database = class_methods(APP_ROOT / "database.py", "ReviewDatabase")
        remaining_http = class_methods(
            APP_ROOT / "backend" / "application.py", "RequestHandler"
        )
        self.assertTrue(HEAVEN_SERVICE_METHODS.isdisjoint(remaining_service))
        self.assertTrue(HEAVEN_REPOSITORY_METHODS.isdisjoint(remaining_database))
        self.assertTrue(HEAVEN_HTTP_METHODS.isdisjoint(remaining_http))

    def test_http_mixin_preserves_all_heaven_endpoints(self) -> None:
        methods = class_methods(
            APP_ROOT / "backend" / "features" / "heaven" / "http.py",
            "HeavenHttpMixin",
        )
        self.assertEqual(set(methods), HEAVEN_HTTP_METHODS)
        source = (
            APP_ROOT / "backend" / "features" / "heaven" / "http.py"
        ).read_text(encoding="utf-8")
        self.assertNotIn("SERVICE.", source)
        self.assertEqual(source.count("self.application_service.heaven_"), 3)


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