diff --git a/tests/unit/test_layer.py b/tests/unit/test_layer.py index 76e048cdf42b8707e719e23dc2f1fde66dec9720..a02552badaca345a9c6349dc5fe6d2f1dfd631ec 100644 --- a/tests/unit/test_layer.py +++ b/tests/unit/test_layer.py @@ -186,13 +186,23 @@ def get_test_indexes(): return indexes +def get_test_encoders(): + encoders = [OpenAIEncoder] + if importlib.util.find_spec("cohere") is not None: + encoders.append(CohereEncoder) + return encoders -@pytest.mark.parametrize("index_cls", get_test_indexes()) -class TestRouteLayer: - def test_initialization(self, openai_encoder, routes, index_cls): + +@pytest.mark.parametrize("index_cls,encoder_cls", [ + (index, encoder) + for index in get_test_indexes() + for encoder in get_test_encoders() +]) +class TestIndexEncoders: + def test_initialization(self, routes, index_cls, encoder_cls): index = init_index(index_cls) route_layer = RouteLayer( - encoder=openai_encoder, routes=routes, index=index, + encoder=encoder_cls(), routes=routes, index=index, auto_sync="local" if index_cls is PineconeIndex else None, top_k=10, ) @@ -210,35 +220,32 @@ class TestRouteLayer: ) def test_initialization_different_encoders( - self, cohere_encoder, openai_encoder, index_cls + self, encoder_cls, index_cls ): index = init_index(index_cls) - route_layer_cohere = RouteLayer(encoder=cohere_encoder, index=index) - assert cohere_encoder.score_threshold == 0.3 - assert route_layer_cohere.score_threshold == 0.3 - route_layer_openai = RouteLayer(encoder=openai_encoder, index=index) - assert route_layer_openai.score_threshold == 0.3 + encoder = encoder_cls() + route_layer = RouteLayer(encoder=encoder, index=index) + assert route_layer.score_threshold == encoder.score_threshold - def test_initialization_no_encoder(self, openai_encoder, index_cls): + def test_initialization_no_encoder(self, openai_encoder, index_cls, encoder_cls): os.environ["OPENAI_API_KEY"] = "test_api_key" route_layer_none = RouteLayer(encoder=None) assert route_layer_none.score_threshold == openai_encoder.score_threshold def test_initialization_dynamic_route( - self, cohere_encoder, openai_encoder, dynamic_routes, index_cls + self, encoder_cls, dynamic_routes, index_cls ): index = init_index(index_cls) - route_layer_cohere = RouteLayer( - encoder=cohere_encoder, routes=dynamic_routes, index=index, + encoder = encoder_cls() + route_layer = RouteLayer( + encoder=encoder, routes=dynamic_routes, index=index, auto_sync="local" if index_cls is PineconeIndex else None, ) - assert route_layer_cohere.score_threshold == 0.3 - route_layer_openai = RouteLayer( - encoder=openai_encoder, routes=dynamic_routes, index=index - ) - assert openai_encoder.score_threshold == 0.3 - assert route_layer_openai.score_threshold == 0.3 + assert route_layer.score_threshold == encoder.score_threshold + +@pytest.mark.parametrize("index_cls", get_test_indexes()) +class TestRouteLayer: def test_delete_index(self, openai_encoder, routes, index_cls): # TODO merge .delete_index() and .delete_all() and get working index = init_index(index_cls) @@ -759,7 +766,10 @@ class TestRouteLayer: def test_retrieve_one_match(self, openai_encoder, routes_3, index_cls): index = init_index(index_cls) - route_layer = RouteLayer(encoder=openai_encoder, routes=routes_3, index=index) + route_layer = RouteLayer( + encoder=openai_encoder, routes=routes_3, index=index, + auto_sync="local", + ) text = "Hello" if index_cls is PineconeIndex: time.sleep(PINECONE_SLEEP) @@ -772,7 +782,10 @@ class TestRouteLayer: self, openai_encoder, routes_2, index_cls ): index = init_index(index_cls) - route_layer = RouteLayer(encoder=openai_encoder, routes=routes_2, index=index) + route_layer = RouteLayer( + encoder=openai_encoder, routes=routes_2, index=index, + auto_sync="local", + ) text = "Hello" if index_cls is PineconeIndex: time.sleep(PINECONE_SLEEP)