tests/test_pinecone_vector_store.py

import importlib
import sys
import tempfile
import types
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

from api.core.models import PineconeCollectionAdapter


def _matches_filter(metadata, filter):
    if not filter:
        return True
    for key, value in filter.items():
        actual = metadata.get(key)
        if isinstance(value, dict) and "$in" in value:
            if actual not in value["$in"]:
                return False
        elif actual != value:
            return False
    return True


class FakeIndex:
    def __init__(self):
        self.vectors = {}

    def upsert(self, vectors, namespace):
        for vector in vectors:
            self.vectors[vector["id"]] = vector

    def fetch(self, ids, namespace):
        return SimpleNamespace(
            vectors={id_: self.vectors[id_] for id_ in ids if id_ in self.vectors}
        )

    def delete(self, ids, namespace):
        for id_ in ids:
            self.vectors.pop(id_, None)

    def query(self, vector, top_k, namespace, filter=None, include_metadata=True):
        matches = []
        for id_, stored in self.vectors.items():
            metadata = stored["metadata"]
            if not _matches_filter(metadata, filter):
                continue
            score = sum(a * b for a, b in zip(vector, stored["values"]))
            matches.append(SimpleNamespace(id=id_, score=score, metadata=metadata))
        matches.sort(key=lambda match: match.score, reverse=True)
        return SimpleNamespace(matches=matches[:top_k])

    def describe_index_stats(self):
        return {"namespaces": {"test": {"vector_count": len(self.vectors)}}}


class PineconeVectorStoreTests(unittest.TestCase):
    def test_pinecone_adapter_preserves_collection_contract(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            adapter = PineconeCollectionAdapter(
                index=FakeIndex(),
                namespace="test",
                cache_path=str(Path(tmp_dir) / "chunks.json"),
            )

            adapter.add(
                ids=["notice-1_0", "notice-2_0"],
                embeddings=[[1.0, 0.0], [0.0, 1.0]],
                documents=["장학금 신청 안내", "인턴십 모집 안내"],
                metadatas=[
                    {"title": "장학금", "url": "https://example.com/1", "category": "장학금"},
                    {"title": "인턴십", "url": "https://example.com/2", "category": "인턴십"},
                ],
            )

            all_data = adapter.get(include=["documents", "metadatas"])
            self.assertEqual(all_data["ids"], ["notice-1_0", "notice-2_0"])
            self.assertEqual(all_data["documents"], ["장학금 신청 안내", "인턴십 모집 안내"])
            self.assertEqual(
                all_data["metadatas"][0],
                {
                    "title": "장학금",
                    "url": "https://example.com/1",
                    "category": "장학금",
                },
            )

            filtered = adapter.get(include=["documents"], where={"category": "인턴십"})
            self.assertEqual(filtered["ids"], ["notice-2_0"])
            self.assertEqual(filtered["documents"], ["인턴십 모집 안내"])

            filtered_in = adapter.get(include=["documents"], where={"category": {"$in": ["장학금", "인턴십"]}})
            self.assertEqual(filtered_in["ids"], ["notice-1_0", "notice-2_0"])
            self.assertEqual(filtered_in["documents"], ["장학금 신청 안내", "인턴십 모집 안내"])

            fetched = adapter.get(ids=["notice-1_0", "missing"], include=["documents", "metadatas"])
            self.assertEqual(fetched["ids"], ["notice-1_0"])
            self.assertEqual(fetched["documents"], ["장학금 신청 안내"])
            self.assertNotIn("document", fetched["metadatas"][0])

            queried = adapter.query(
                query_embeddings=[[0.9, 0.1]],
                n_results=1,
                include=["metadatas", "distances", "documents"],
            )
            self.assertEqual(queried["ids"], [["notice-1_0"]])
            self.assertEqual(
                queried["metadatas"],
                [[{
                    "title": "장학금",
                    "url": "https://example.com/1",
                    "category": "장학금",
                }]],
            )
            self.assertEqual(queried["documents"], [["장학금 신청 안내"]])
            self.assertAlmostEqual(queried["distances"][0][0], 0.1)

            adapter.delete(ids=["notice-1_0"])
            self.assertEqual(adapter.get(ids=["notice-1_0"])["ids"], [])
            self.assertEqual(adapter.count(), 1)

    def test_get_vector_collection_can_initialize_pinecone_adapter(self):
        fake_index = FakeIndex()

        class FakePinecone:
            def __init__(self, api_key):
                self.api_key = api_key
                self.created = []

            def has_index(self, name):
                return False

            def create_index(self, **kwargs):
                self.created.append(kwargs)

            def describe_index(self, name):
                return {"status": {"ready": True}, "dimension": 2}

            def Index(self, name):
                return fake_index

        fake_module = types.ModuleType("pinecone")
        fake_module.Pinecone = FakePinecone
        fake_module.ServerlessSpec = lambda cloud, region: {"cloud": cloud, "region": region}

        with tempfile.TemporaryDirectory() as tmp_dir:
            models = importlib.import_module("api.core.models")
            with (
                mock.patch.dict(sys.modules, {"pinecone": fake_module}),
                mock.patch.object(models, "_vector_collection", None),
                mock.patch.object(models, "VECTOR_DB", "pinecone"),
                mock.patch.object(models, "PINECONE_API_KEY", "test-key"),
                mock.patch.object(models, "PINECONE_INDEX_NAME", "hansung-notices-test"),
                mock.patch.object(models, "PINECONE_NAMESPACE", "test"),
                mock.patch.object(models, "PINECONE_CACHE_PATH", str(Path(tmp_dir) / "chunks.json")),
                mock.patch.object(models, "EMBEDDING_DIM", 2),
            ):
                collection = models.get_vector_collection()

        self.assertIsInstance(collection, PineconeCollectionAdapter)
        self.assertIs(collection.index, fake_index)
        self.assertEqual(collection.namespace, "test")


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