boost-data-collector/core/tests/adapters/test_pinecone_adapter.py at develop · cppalliance/boost-data-collector · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
"""Unit tests for core.adapters.pinecone.PineconeAdapter."""

from unittest.mock import MagicMock, patch

import pytest

from core.adapters.pinecone import (
    PineconeAdapter,
    PineconeSdkIndexAdapter,
    ensure_pinecone_available,
)


pytest.importorskip("pinecone")


def test_ensure_pinecone_available_when_installed():
    ensure_pinecone_available()


def test_pinecone_adapter_from_api_key():
    with patch("core.adapters.pinecone.Pinecone") as mock_pc_cls:
        mock_pc = MagicMock()
        mock_pc_cls.return_value = mock_pc
        adapter = PineconeAdapter.from_api_key("pc-test")
    mock_pc_cls.assert_called_once_with(api_key="pc-test")
    assert adapter._pc is mock_pc


def test_pinecone_adapter_list_index_names():
    idx_a = MagicMock()
    idx_a.name = "dense"
    idx_b = MagicMock()
    idx_b.name = "dense-sparse"
    pc = MagicMock()
    pc.list_indexes.return_value = [idx_a, idx_b]
    adapter = PineconeAdapter(pc)
    assert adapter.list_index_names() == {"dense", "dense-sparse"}


def test_pinecone_adapter_create_index_for_model():
    pc = MagicMock()
    adapter = PineconeAdapter(pc)
    embed = {"model": "m", "field_map": {"text": "chunk_text"}}
    adapter.create_index_for_model(
        name="idx",
        cloud="aws",
        region="us-east-1",
        embed=embed,
    )
    pc.create_index_for_model.assert_called_once_with(
        name="idx",
        cloud="aws",
        region="us-east-1",
        embed=embed,
    )


def test_pinecone_sdk_index_adapter_forwards_calls():
    raw = MagicMock()
    idx = PineconeSdkIndexAdapter(raw)
    idx.upsert_records(records=[{"id": "1"}], namespace="ns")
    raw.upsert_records.assert_called_once_with(
        records=[{"id": "1"}],
        namespace="ns",
    )
    idx.update(id="1", set_metadata={"k": "v"}, namespace="ns")
    raw.update.assert_called_once_with(
        id="1",
        set_metadata={"k": "v"},
        namespace="ns",
    )
    idx.delete(ids=["1"], namespace="ns")
    raw.delete.assert_called_once_with(ids=["1"], namespace="ns")
    idx.describe_index_stats()
    raw.describe_index_stats.assert_called_once()