Skip to content

Commit

Permalink
Fix unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@jhamon
jhamon committed Dec 17, 2024
1 parent 9d8e5b8 commit 534df9b
Showing 1 changed file with 73 additions and 66 deletions.
139 changes: 73 additions & 66 deletions tests/unit/test_index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
import pytest

from pinecone.data import _Index
import pinecone.core.openapi.db_data.models as oai_models
from pinecone import QueryResponse, UpsertResponse, SparseValues, Vector
import pinecone.core.openapi.db_data.models as oai
from pinecone import QueryResponse, UpsertResponse, Vector


class TestRestIndex:
Expand Down Expand Up @@ -32,10 +32,10 @@ def test_upsert_tuplesOfIdVec_UpserWithoutMD(self, mocker):
mocker.patch.object(self.index._vector_api, "upsert_vectors", autospec=True)
self.index.upsert([("vec1", self.vals1), ("vec2", self.vals2)], namespace="ns")
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata={}),
Vector(id="vec2", values=self.vals2, metadata={}),
oai.Vector(id="vec1", values=self.vals1, metadata={}),
oai.Vector(id="vec2", values=self.vals2, metadata={}),
],
namespace="ns",
)
Expand All @@ -45,10 +45,10 @@ def test_upsert_tuplesOfIdVecMD_UpsertVectorsWithMD(self, mocker):
mocker.patch.object(self.index._vector_api, "upsert_vectors", autospec=True)
self.index.upsert([("vec1", self.vals1, self.md1), ("vec2", self.vals2, self.md2)])
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
]
)
)
Expand All @@ -62,10 +62,10 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithMD(self, mocker):
]
)
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
]
)
)
Expand All @@ -76,8 +76,11 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithoutMD(self, mocker):
[{"id": self.id1, "values": self.vals1}, {"id": self.id2, "values": self.vals2}]
)
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
vectors=[Vector(id="vec1", values=self.vals1), Vector(id="vec2", values=self.vals2)]
oai.UpsertRequest(
vectors=[
oai.Vector(id="vec1", values=self.vals1),
oai.Vector(id="vec2", values=self.vals2),
]
)
)

Expand All @@ -90,10 +93,14 @@ def test_upsert_dictOfIdVecMD_UpsertVectorsWithSparseValues(self, mocker):
]
)
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, sparse_values=SparseValues(**self.sv1)),
Vector(id="vec2", values=self.vals2, sparse_values=SparseValues(**self.sv2)),
oai.Vector(
id="vec1", values=self.vals1, sparse_values=oai.SparseValues(**self.sv1)
),
oai.Vector(
id="vec2", values=self.vals2, sparse_values=oai.SparseValues(**self.sv2)
),
]
)
)
Expand All @@ -102,16 +109,16 @@ def test_upsert_vectors_upsertInputVectors(self, mocker):
mocker.patch.object(self.index._vector_api, "upsert_vectors", autospec=True)
self.index.upsert(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
],
namespace="ns",
)
self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
],
namespace="ns",
)
Expand All @@ -134,16 +141,16 @@ def test_upsert_parallelUpsert_callUpsertParallel(self, mocker):
[async_result.get() for async_result in async_results]

index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec1", values=self.vals1, metadata=self.md1)],
oai.UpsertRequest(
vectors=[oai.Vector(id="vec1", values=self.vals1, metadata=self.md1)],
namespace="ns",
),
async_req=True,
)

index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec2", values=self.vals2, metadata=self.md2)],
oai.UpsertRequest(
vectors=[oai.Vector(id="vec2", values=self.vals2, metadata=self.md2)],
namespace="ns",
),
async_req=True,
Expand All @@ -170,14 +177,16 @@ def test_upsert_vectorListIsMultiplyOfBatchSize_vectorsUpsertedInBatches(self, m
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec1", values=self.vals1, metadata=self.md1)], namespace="ns"
oai.UpsertRequest(
vectors=[oai.Vector(id="vec1", values=self.vals1, metadata=self.md1)],
namespace="ns",
)
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec2", values=self.vals2, metadata=self.md2)], namespace="ns"
oai.UpsertRequest(
vectors=[oai.Vector(id="vec2", values=self.vals2, metadata=self.md2)],
namespace="ns",
)
)

Expand All @@ -195,27 +204,28 @@ def test_upsert_vectorListNotMultiplyOfBatchSize_vectorsUpsertedInBatches(self,

result = self.index.upsert(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
Vector(id="vec3", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec3", values=self.vals1, metadata=self.md1),
],
namespace="ns",
batch_size=2,
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
],
namespace="ns",
)
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec3", values=self.vals1, metadata=self.md1)], namespace="ns"
oai.UpsertRequest(
vectors=[oai.Vector(id="vec3", values=self.vals1, metadata=self.md1)],
namespace="ns",
)
)

Expand All @@ -242,11 +252,11 @@ def test_upsert_vectorListSmallerThanBatchSize_vectorsUpsertedInBatches(self, mo
)

self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
Vector(id="vec3", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec3", values=self.vals1, metadata=self.md1),
],
namespace="ns",
)
Expand Down Expand Up @@ -275,18 +285,19 @@ def test_upsert_tuplesList_vectorsUpsertedInBatches(self, mocker):
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
],
namespace="ns",
)
)

self.index._vector_api.upsert_vectors.assert_any_call(
oai_models.UpsertRequest(
vectors=[Vector(id="vec3", values=self.vals1, metadata=self.md1)], namespace="ns"
oai.UpsertRequest(
vectors=[oai.Vector(id="vec3", values=self.vals1, metadata=self.md1)],
namespace="ns",
)
)

Expand All @@ -308,10 +319,10 @@ def test_upsert_dataframe(self, mocker):
self.index.upsert_from_dataframe(df)

self.index._vector_api.upsert_vectors.assert_called_once_with(
oai_models.UpsertRequest(
oai.UpsertRequest(
vectors=[
Vector(id="vec1", values=self.vals1, metadata=self.md1),
Vector(id="vec2", values=self.vals2, metadata=self.md2),
oai.Vector(id="vec1", values=self.vals1, metadata=self.md1),
oai.Vector(id="vec2", values=self.vals2, metadata=self.md2),
]
)
)
Expand All @@ -326,7 +337,7 @@ def test_upsert_batchSizeIsNotPositive_errorIsRaised(self):

with pytest.raises(ValueError, match="batch_size must be a positive integer"):
self.index.upsert(
vectors=[Vector(id="vec1", values=self.vals1, metadata=self.md1)],
vectors=[oai.Vector(id="vec1", values=self.vals1, metadata=self.md1)],
namespace="ns",
batch_size=-1,
)
Expand All @@ -349,7 +360,7 @@ def test_upsert_useBatchSizeAndAsyncReq_valueErrorRaised(self):
def test_query_byVectorNoFilter_queryVectorNoFilter(self, mocker):
response = QueryResponse(
results=[],
matches=[oai_models.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})],
matches=[oai.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})],
namespace="test",
)
mocker.patch.object(
Expand All @@ -359,10 +370,10 @@ def test_query_byVectorNoFilter_queryVectorNoFilter(self, mocker):
actual = self.index.query(top_k=10, vector=self.vals1)

self.index._vector_api.query_vectors.assert_called_once_with(
oai_models.QueryRequest(top_k=10, vector=self.vals1)
oai.QueryRequest(top_k=10, vector=self.vals1)
)
expected = QueryResponse(
matches=[oai_models.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})],
matches=[oai.ScoredVector(id="1", score=0.9, values=[0.0], metadata={"a": 2})],
namespace="test",
)
assert expected.to_dict() == actual.to_dict()
Expand All @@ -371,18 +382,14 @@ def test_query_byVectorWithFilter_queryVectorWithFilter(self, mocker):
mocker.patch.object(self.index._vector_api, "query_vectors", autospec=True)
self.index.query(top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns")
self.index._vector_api.query_vectors.assert_called_once_with(
oai_models.QueryRequest(
top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns"
)
oai.QueryRequest(top_k=10, vector=self.vals1, filter=self.filter1, namespace="ns")
)

def test_query_byVecId_queryByVecId(self, mocker):
mocker.patch.object(self.index._vector_api, "query_vectors", autospec=True)
self.index.query(top_k=10, id="vec1", include_metadata=True, include_values=False)
self.index._vector_api.query_vectors.assert_called_once_with(
oai_models.QueryRequest(
top_k=10, id="vec1", include_metadata=True, include_values=False
)
oai.QueryRequest(top_k=10, id="vec1", include_metadata=True, include_values=False)
)

def test_query_rejects_both_id_and_vector(self):
Expand All @@ -405,21 +412,21 @@ def test_delete_byIds_deleteByIds(self, mocker):
mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True)
self.index.delete(ids=["vec1", "vec2"])
self.index._vector_api.delete_vectors.assert_called_once_with(
oai_models.DeleteRequest(ids=["vec1", "vec2"])
oai.DeleteRequest(ids=["vec1", "vec2"])
)

def test_delete_deleteAllByFilter_deleteAllByFilter(self, mocker):
mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True)
self.index.delete(delete_all=True, filter=self.filter1, namespace="ns")
self.index._vector_api.delete_vectors.assert_called_once_with(
oai_models.DeleteRequest(delete_all=True, filter=self.filter1, namespace="ns")
oai.DeleteRequest(delete_all=True, filter=self.filter1, namespace="ns")
)

def test_delete_deleteAllNoFilter_deleteNoFilter(self, mocker):
mocker.patch.object(self.index._vector_api, "delete_vectors", autospec=True)
self.index.delete(delete_all=True)
self.index._vector_api.delete_vectors.assert_called_once_with(
oai_models.DeleteRequest(delete_all=True)
oai.DeleteRequest(delete_all=True)
)

# endregion
Expand All @@ -446,14 +453,14 @@ def test_update_byIdAnValues_updateByIdAndValues(self, mocker):
mocker.patch.object(self.index._vector_api, "update_vector", autospec=True)
self.index.update(id="vec1", values=self.vals1, namespace="ns")
self.index._vector_api.update_vector.assert_called_once_with(
oai_models.UpdateRequest(id="vec1", values=self.vals1, namespace="ns")
oai.UpdateRequest(id="vec1", values=self.vals1, namespace="ns")
)

def test_update_byIdAnValuesAndMetadata_updateByIdAndValuesAndMetadata(self, mocker):
mocker.patch.object(self.index._vector_api, "update_vector", autospec=True)
self.index.update("vec1", values=self.vals1, metadata=self.md1)
self.index._vector_api.update_vector.assert_called_once_with(
oai_models.UpdateRequest(id="vec1", values=self.vals1, metadata=self.md1)
oai.UpdateRequest(id="vec1", values=self.vals1, metadata=self.md1)
)

# endregion
Expand All @@ -464,14 +471,14 @@ def test_describeIndexStats_callWithoutFilter_CalledWithoutFilter(self, mocker):
mocker.patch.object(self.index._vector_api, "describe_index_stats", autospec=True)
self.index.describe_index_stats()
self.index._vector_api.describe_index_stats.assert_called_once_with(
oai_models.DescribeIndexStatsRequest()
oai.DescribeIndexStatsRequest()
)

def test_describeIndexStats_callWithFilter_CalledWithFilter(self, mocker):
mocker.patch.object(self.index._vector_api, "describe_index_stats", autospec=True)
self.index.describe_index_stats(filter=self.filter1)
self.index._vector_api.describe_index_stats.assert_called_once_with(
oai_models.DescribeIndexStatsRequest(filter=self.filter1)
oai.DescribeIndexStatsRequest(filter=self.filter1)
)

# endregion

0 comments on commit 534df9b

Please sign in to comment.