replypier/backend/tests/test_enterprise_optional_import_contract.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import ast
import builtins
from contextlib import redirect_stdout
from dataclasses import dataclass
from io import StringIO
from pathlib import Path
from types import ModuleType
from typing import Literal

import pytest


BACKEND_ROOT = Path(__file__).resolve().parents[1]


@dataclass(frozen=True)
class EnterpriseProbe:
    source: str
    module: str
    symbol: str
    occurrence: int = 0


PROBES = (
    EnterpriseProbe(
        "app/api/agent.py",
        "app.enterprise.repositories.subscription",
        "SubscriptionRepository",
        occurrence=0,
    ),
    EnterpriseProbe(
        "app/api/agent.py",
        "app.enterprise.repositories.subscription",
        "SubscriptionRepository",
        occurrence=1,
    ),
    EnterpriseProbe(
        "app/knowledge/knowledge_base.py",
        "app.enterprise.repositories.subscription",
        "SubscriptionRepository",
    ),
    EnterpriseProbe("app/main.py", "app.enterprise.models", "OTP"),
)


def _probe_node(probe: EnterpriseProbe) -> ast.Try:
    source_path = BACKEND_ROOT / probe.source
    tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path))
    matches = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Try):
            continue
        imports_symbol = any(
            isinstance(statement, ast.ImportFrom)
            and statement.module == probe.module
            and any(imported.name == probe.symbol for imported in statement.names)
            for statement in node.body
        )
        if imports_symbol:
            matches.append(node)
    matches.sort(key=lambda node: node.lineno)
    assert len(matches) > probe.occurrence, (
        f"{probe.source} no longer contains the requested optional import probe"
    )
    return matches[probe.occurrence]


def _execute_probe(
    probe: EnterpriseProbe,
    outcome: Literal["missing_module", "missing_symbol", "available", "internal_attribute_error"],
) -> tuple[dict[str, object], str]:
    probe_node = _probe_node(probe)
    source_path = BACKEND_ROOT / probe.source
    snippet = ast.Module(body=[probe_node], type_ignores=[])
    code = compile(snippet, str(source_path), "exec")
    available_target = type("SyntheticEnterpriseTarget", (), {})

    def fake_import(
        name: str,
        globals: dict[str, object] | None = None,
        locals: dict[str, object] | None = None,
        fromlist: tuple[str, ...] = (),
        level: int = 0,
    ) -> object:
        del globals, locals, fromlist, level
        if name.startswith("app.enterprise"):
            if outcome == "missing_module":
                raise ModuleNotFoundError(name)
            if outcome == "internal_attribute_error":
                raise AttributeError("synthetic enterprise import-time failure")
            module = ModuleType(name)
            if outcome == "available":
                setattr(module, probe.symbol, available_target)
            return module
        return builtins.__import__(name)

    namespace: dict[str, object] = {
        "__builtins__": {**vars(builtins), "__import__": fake_import},
        "__name__": "enterprise_probe_contract_test",
    }
    output = StringIO()
    with redirect_stdout(output):
        exec(code, namespace)
    return namespace, output.getvalue()


@pytest.mark.parametrize("probe", PROBES, ids=lambda probe: f"{probe.source}:{probe.occurrence}")
@pytest.mark.parametrize("outcome", ["missing_module", "missing_symbol"])
def test_optional_probe_falls_back_when_module_or_symbol_is_missing(
    probe: EnterpriseProbe,
    outcome: Literal["missing_module", "missing_symbol"],
) -> None:
    namespace, output = _execute_probe(probe, outcome)

    if "HAS_ENTERPRISE" in namespace:
        assert namespace["HAS_ENTERPRISE"] is False
    else:
        assert output == "Enterprise models not available\n"


@pytest.mark.parametrize("probe", PROBES, ids=lambda probe: f"{probe.source}:{probe.occurrence}")
def test_optional_probe_accepts_a_valid_enterprise_symbol(probe: EnterpriseProbe) -> None:
    namespace, output = _execute_probe(probe, "available")

    assert namespace[probe.symbol].__name__ == "SyntheticEnterpriseTarget"
    if "HAS_ENTERPRISE" in namespace:
        assert namespace["HAS_ENTERPRISE"] is True
    else:
        assert output == ""


@pytest.mark.parametrize("probe", PROBES, ids=lambda probe: f"{probe.source}:{probe.occurrence}")
def test_import_time_attribute_error_is_not_mistaken_for_absent_enterprise(
    probe: EnterpriseProbe,
) -> None:
    with pytest.raises(AttributeError, match="synthetic enterprise import-time failure"):
        _execute_probe(probe, "internal_attribute_error")