xiaobaifupan/app/tests/test_preservation_slice_mentor_llm.py

from __future__ import annotations

import ast
import hashlib
import unittest
from pathlib import Path

import llm_stream
import mentor_agent
from backend.features.mentor import agent as canonical_agent
from backend.llm import stream as canonical_stream


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

MENTOR_SERVICE_METHODS = {
    "mentor_setup",
    "save_mentor_preferences",
    "mentor_stream",
    "mentor_messages",
    "clear_mentor_messages",
    "_validate_mentor_history",
    "_build_mentor_context",
    "_mentor_market_matrix",
}

LLM_SERVICE_METHODS = {
    "_personal_llm_profile",
    "_platform_llm_profile",
    "_profile_configured",
    "_resolved_llm_profile",
    "llm_primary_api_key",
    "llm_primary_base_url",
    "llm_primary_model",
    "llm_fallback_api_key",
    "llm_fallback_base_url",
    "llm_fallback_model",
    "llm_source",
    "llm_configured",
    "llm_fallback_configured",
    "save_llm_settings",
    "save_llm_mode",
    "test_llm_profile",
    "_validate_llm_profile",
    "llm_access_status",
    "_platform_usage_today",
    "_platform_usage_today_for_user",
    "test_system_llm_profile",
}

MENTOR_REPOSITORY_METHODS = {
    "save_mentor_exchange",
    "list_mentor_messages",
    "delete_mentor_messages",
    "list_mentor_preferences",
    "save_mentor_preferences",
}

LLM_REPOSITORY_METHODS = {"record_llm_usage", "count_llm_usage_since"}


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 assignments(path: Path, names: set[str]) -> dict[str, str]:
    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
    result = {}
    for node in tree.body:
        if not isinstance(node, ast.Assign) or len(node.targets) != 1:
            continue
        target = node.targets[0]
        if isinstance(target, ast.Name) and target.id in names:
            result[target.id] = ast.dump(node.value, include_attributes=False)
    return result


class MentorLLMSliceSourceEquivalenceTests(unittest.TestCase):
    def assert_methods_equal(
        self,
        original_path: Path,
        original_class: str,
        migrated_path: Path,
        migrated_class: str,
        expected: set[str],
    ) -> 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):
            self.assertEqual(migrated[name], original[name], name)

    def test_mentor_agent_and_stream_accumulator_are_exact_files(self) -> None:
        self.assertEqual(
            sha256(ORIGINAL_ROOT / "mentor_agent.py"),
            sha256(APP_ROOT / "backend" / "features" / "mentor" / "agent.py"),
        )
        self.assertEqual(
            sha256(ORIGINAL_ROOT / "llm_stream.py"),
            sha256(APP_ROOT / "backend" / "llm" / "stream.py"),
        )

    def test_compatibility_modules_are_canonical_module_objects(self) -> None:
        self.assertIs(mentor_agent, canonical_agent)
        self.assertIs(llm_stream, canonical_stream)

    def test_mentor_service_methods_are_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "server.py",
            "DashboardService",
            APP_ROOT / "backend" / "features" / "mentor" / "service.py",
            "MentorServiceMixin",
            MENTOR_SERVICE_METHODS,
        )

    def test_llm_service_methods_are_exact_original_ast(self) -> None:
        original = class_methods(ORIGINAL_ROOT / "server.py", "DashboardService")
        migrated = class_methods(
            APP_ROOT / "backend" / "llm" / "service.py", "LLMServiceMixin"
        )
        self.assertEqual(set(migrated), LLM_SERVICE_METHODS)
        adapted = {"_platform_usage_today", "_platform_usage_today_for_user"}
        for name in sorted(LLM_SERVICE_METHODS - adapted):
            self.assertEqual(migrated[name], original[name], name)
        source = (APP_ROOT / "backend" / "llm" / "service.py").read_text(
            encoding="utf-8"
        )
        self.assertIn(
            "return self._platform_usage_today_for_user(self.current_user_id)", source
        )
        self.assertIn("def _platform_usage_today_for_user(self, user_id: int)", source)

    def test_mentor_and_llm_repositories_are_exact_original_ast(self) -> None:
        self.assert_methods_equal(
            ORIGINAL_ROOT / "database.py",
            "ReviewDatabase",
            APP_ROOT / "backend" / "features" / "mentor" / "repository.py",
            "MentorRepositoryMixin",
            MENTOR_REPOSITORY_METHODS,
        )
        self.assert_methods_equal(
            ORIGINAL_ROOT / "database.py",
            "ReviewDatabase",
            APP_ROOT / "backend" / "llm" / "repository.py",
            "LLMAuditRepositoryMixin",
            LLM_REPOSITORY_METHODS,
        )

    def test_mentor_data_profiles_are_exact_original_values(self) -> None:
        names = {"MENTOR_DATA_PROFILES", "MENTOR_INDEX_UNIVERSE", "MENTOR_ETF_UNIVERSE"}
        self.assertEqual(
            assignments(ORIGINAL_ROOT / "server.py", names),
            assignments(
                APP_ROOT / "backend" / "features" / "mentor" / "service.py",
                names,
            ),
        )

    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(MENTOR_SERVICE_METHODS.isdisjoint(remaining_service))
        self.assertTrue(LLM_SERVICE_METHODS.isdisjoint(remaining_service))
        self.assertTrue(MENTOR_REPOSITORY_METHODS.isdisjoint(remaining_database))
        self.assertTrue(LLM_REPOSITORY_METHODS.isdisjoint(remaining_database))
        self.assertTrue(
            {"stream_mentor_chat", "save_llm_settings", "save_llm_mode", "test_llm_settings"}
            .isdisjoint(remaining_http)
        )

    def test_http_mixins_preserve_stream_and_model_endpoints(self) -> None:
        mentor_http = class_methods(
            APP_ROOT / "backend" / "features" / "mentor" / "http.py",
            "MentorHttpMixin",
        )
        llm_http = class_methods(APP_ROOT / "backend" / "llm" / "http.py", "LLMHttpMixin")
        self.assertEqual(set(mentor_http), {"stream_mentor_chat"})
        self.assertEqual(
            set(llm_http), {"save_llm_settings", "save_llm_mode", "test_llm_settings"}
        )


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