From 534df9bb67c7c358b8329d91b4217bf7caa84018 Mon Sep 17 00:00:00 2001 From: Jen Hamon Date: Mon, 16 Dec 2024 20:16:13 -0500 Subject: [PATCH] Fix unit tests --- tests/unit/test_index.py | 139 ++++++++++++++++++++------------------- 1 file changed, 73 insertions(+), 66 deletions(-) diff --git a/tests/unit/test_index.py b/tests/unit/test_index.py index fe45108c..72ed7422 100644 --- a/tests/unit/test_index.py +++ b/tests/unit/test_index.py @@ -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: @@ -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", ) @@ -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), ] ) ) @@ -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), ] ) ) @@ -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), + ] ) ) @@ -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) + ), ] ) ) @@ -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", ) @@ -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, @@ -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", ) ) @@ -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", ) ) @@ -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", ) @@ -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", ) ) @@ -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), ] ) ) @@ -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, ) @@ -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( @@ -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() @@ -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): @@ -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 @@ -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 @@ -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