haystack-core-integrations/integrations/pgvector/tests/conftest.py at 092e37d7746ca4466f778edd5bb723bd13ae8a33 · deepset-ai/haystack-core-integrations · 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
from unittest.mock import Mock, patch

import pytest
from psycopg import AsyncConnection, Connection

from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore


@pytest.fixture
def document_store(request, monkeypatch):
    monkeypatch.setenv("PG_CONN_STR", "postgresql://postgres:postgres@localhost:5432/postgres")
    table_name = f"haystack_{request.node.name}"
    embedding_dimension = 768
    vector_function = "cosine_similarity"
    recreate_table = True
    search_strategy = "exact_nearest_neighbor"

    store = PgvectorDocumentStore(
        table_name=table_name,
        embedding_dimension=embedding_dimension,
        vector_function=vector_function,
        recreate_table=recreate_table,
        search_strategy=search_strategy,
    )

    yield store

    # Ensure connection just for table deletion.
    # During test execution, the tested methods are expected to call _ensure_db_setup() themselves.
    store._ensure_db_setup()
    store.delete_table()


@pytest.fixture
def document_store_w_hnsw_index(request, monkeypatch):
    monkeypatch.setenv("PG_CONN_STR", "postgresql://postgres:postgres@localhost:5432/postgres")
    table_name = f"haystack_hnsw_{request.node.name}"
    embedding_dimension = 768
    vector_function = "cosine_similarity"
    recreate_table = True
    search_strategy = "hnsw"

    store = PgvectorDocumentStore(
        table_name=table_name,
        embedding_dimension=embedding_dimension,
        vector_function=vector_function,
        recreate_table=recreate_table,
        search_strategy=search_strategy,
    )
    yield store

    # Ensure connection just for table deletion.
    # During test execution, the tested methods are expected to call _ensure_db_setup() themselves.
    store._ensure_db_setup()
    store.delete_table()


@pytest.fixture
def document_store_w_halfvec_hnsw_index(request, monkeypatch):
    monkeypatch.setenv("PG_CONN_STR", "postgresql://postgres:postgres@localhost:5432/postgres")
    table_name = f"haystack_halfvec_hnsw_{request.node.name}"
    embedding_dimension = 2500
    vector_function = "cosine_similarity"
    recreate_table = True

    store = PgvectorDocumentStore(
        table_name=table_name,
        embedding_dimension=embedding_dimension,
        vector_function=vector_function,
        recreate_table=recreate_table,
        search_strategy="hnsw",
        vector_type="halfvec",
    )
    yield store

    # Ensure connection just for table deletion.
    # During test execution, the tested methods are expected to call _ensure_db_setup() themselves.
    store._ensure_db_setup()
    store.delete_table()


@pytest.fixture
def patches_for_unit_tests():
    with (
        patch("haystack_integrations.document_stores.pgvector.document_store.register_vector") as mock_register,
        patch(
            "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore.delete_table"
        ) as mock_delete,
        patch(
            "haystack_integrations.document_stores.pgvector.document_store.PgvectorDocumentStore._handle_hnsw"
        ) as mock_hnsw,
    ):
        yield mock_register, mock_delete, mock_hnsw


@pytest.fixture
def mock_store(patches_for_unit_tests, monkeypatch):  # noqa: ARG001  patches are not explicitly called but necessary
    monkeypatch.setenv("PG_CONN_STR", "some-connection-string")
    table_name = "haystack"
    embedding_dimension = 768
    vector_function = "cosine_similarity"
    recreate_table = True
    search_strategy = "exact_nearest_neighbor"

    store = PgvectorDocumentStore(
        table_name=table_name,
        embedding_dimension=embedding_dimension,
        vector_function=vector_function,
        recreate_table=recreate_table,
        search_strategy=search_strategy,
    )

    yield store


@pytest.fixture
def mock_store_with_mock_connection(mock_store):
    mock_store._connection = Mock(spec=Connection)
    return mock_store


@pytest.fixture
def mock_store_with_mock_async_connection(mock_store):
    mock_store._async_connection = Mock(spec=AsyncConnection)
    return mock_store