daily-paper/tests/test_embedder.py

"""Embedder / Chroma 服务测试 — 初始化、索引、embedding API。"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from app.config import settings


# ═══════════════════════════════════════════════════════════════════════
# 初始化
# ═══════════════════════════════════════════════════════════════════════


class TestEmbedderInit:
    """embedder.py 初始化测试。"""

    def test_chroma_disabled_skip_init(self, monkeypatch):
        """CHROMA_ENABLED=false 时不初始化。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        emb.init_chroma()
        assert emb._chroma._client is None

    def test_chroma_init_success(self, monkeypatch, tmp_path):
        """CHROMA_ENABLED=true 时初始化成功。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", True)
        monkeypatch.setattr(settings, "CHROMA_DIR", str(tmp_path / "chroma"))

        import app.services.embedder as emb

        emb._chroma.reset()
        emb.init_chroma()

        assert emb._chroma._client is not None
        assert emb._chroma._collection is not None

        # 清理
        emb._chroma.reset()

    def test_get_collection_returns_none_when_disabled(self, monkeypatch):
        """CHROMA_ENABLED=false 时 get_collection 返回 None。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        assert emb.get_collection() is None


# ═══════════════════════════════════════════════════════════════════════
# 索引
# ═══════════════════════════════════════════════════════════════════════


class TestEmbedderIndexing:
    """embedder.py 索引测试。"""

    def test_index_paper_disabled(self, monkeypatch):
        """CHROMA_ENABLED=false 时 index_paper 返回 False。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        assert emb.index_paper("test-id") is False

    def test_index_paper_no_api_config(self, monkeypatch, tmp_path):
        """没有 EMBED_API_BASE 时返回 False。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", True)
        monkeypatch.setattr(settings, "CHROMA_DIR", str(tmp_path / "chroma"))
        monkeypatch.setattr(settings, "EMBED_API_BASE", "")
        monkeypatch.setattr(settings, "EMBED_MODEL", "")

        import app.services.embedder as emb

        emb._chroma.reset()
        emb.init_chroma()

        result = emb.index_paper("test-id", {"title_zh": "测试", "title_en": "Test"})
        assert result is False

        emb._chroma.reset()

    def test_index_batch_disabled(self, monkeypatch):
        """CHROMA_ENABLED=false 时 index_batch 返回全失败。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        result = emb.index_batch(["a", "b"])
        assert result["success"] == 0
        assert result["failed"] == 2

    def test_index_batch_empty(self, monkeypatch):
        """空列表时返回 0。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        result = emb.index_batch([])
        assert result["total"] == 0

    def test_delete_paper_disabled(self, monkeypatch):
        """CHROMA_ENABLED=false 时 delete_paper 返回 False。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        assert emb.delete_paper("test-id") is False

    def test_search_similar_disabled(self, monkeypatch):
        """CHROMA_ENABLED=false 时 search_similar 返回空列表。"""
        monkeypatch.setattr(settings, "CHROMA_ENABLED", False)
        import app.services.embedder as emb

        emb._chroma.reset()
        assert emb.search_similar("test query") == []


# ═══════════════════════════════════════════════════════════════════════
# Embedding API
# ═══════════════════════════════════════════════════════════════════════


class TestEmbeddingApi:
    """_get_embedding 测试。"""

    def test_no_api_base_returns_none(self, monkeypatch):
        """EMBED_API_BASE 为空时返回 None。"""
        monkeypatch.setattr(settings, "EMBED_API_BASE", "")
        monkeypatch.setattr(settings, "EMBED_MODEL", "")
        import app.services.embedder as emb

        assert emb._get_embedding("test") is None

    def test_dimension_mismatch_returns_none(self, monkeypatch):
        """维度不匹配时返回 None。"""
        monkeypatch.setattr(settings, "EMBED_API_BASE", "http://fake")
        monkeypatch.setattr(settings, "EMBED_MODEL", "test-model")
        monkeypatch.setattr(settings, "EMBED_API_KEY", "")
        monkeypatch.setattr(settings, "EMBED_DIMENSIONS", 128)
        monkeypatch.setattr(settings, "HTTP_TIMEOUT_SECONDS", 5)

        import app.services.embedder as emb

        mock_resp = MagicMock()
        mock_resp.json.return_value = {"data": [{"embedding": [0.1] * 64}]}
        mock_resp.raise_for_status = MagicMock()

        with patch("httpx.Client") as mock_client:
            mock_client.return_value.__enter__ = MagicMock(return_value=mock_resp)
            mock_client.return_value.__exit__ = MagicMock(return_value=False)
            result = emb._get_embedding("test")
            assert result is None

    def test_api_failure_returns_none(self, monkeypatch):
        """API 调用失败时返回 None。"""
        monkeypatch.setattr(settings, "EMBED_API_BASE", "http://fake")
        monkeypatch.setattr(settings, "EMBED_MODEL", "test-model")
        monkeypatch.setattr(settings, "EMBED_API_KEY", "")
        monkeypatch.setattr(settings, "EMBED_DIMENSIONS", 0)
        monkeypatch.setattr(settings, "HTTP_TIMEOUT_SECONDS", 5)

        import app.services.embedder as emb

        with patch("httpx.Client") as mock_client:
            mock_client.return_value.__enter__ = MagicMock()
            mock_client.return_value.__exit__ = MagicMock(return_value=False)
            mock_client.return_value.__enter__.return_value.post.side_effect = (
                Exception("timeout")
            )
            result = emb._get_embedding("test")
            assert result is None