fix text embedder tests

deepset-ai · Jan 16, 2024 · 314a293 · 314a293
1 parent a0ce770
commit 314a293
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/integrations/cohere/tests/test_text_embedder.py b/integrations/cohere/tests/test_text_embedder.py
@@ -19,7 +19,7 @@ def test_init_default(self):
         embedder = CohereTextEmbedder(api_key="test-api-key")
 
         assert embedder.api_key == "test-api-key"
-        assert embedder.model_name == "embed-english-v2.0"
+        assert embedder.model == "embed-english-v2.0"
         assert embedder.input_type == "search_query"
         assert embedder.api_base_url == COHERE_API_URL
         assert embedder.truncate == "END"
@@ -33,7 +33,7 @@ def test_init_with_parameters(self):
         """
         embedder = CohereTextEmbedder(
             api_key="test-api-key",
-            model_name="embed-multilingual-v2.0",
+            model="embed-multilingual-v2.0",
             input_type="classification",
             api_base_url="https://custom-api-base-url.com",
             truncate="START",
@@ -42,7 +42,7 @@ def test_init_with_parameters(self):
             timeout=60,
         )
         assert embedder.api_key == "test-api-key"
-        assert embedder.model_name == "embed-multilingual-v2.0"
+        assert embedder.model == "embed-multilingual-v2.0"
         assert embedder.input_type == "classification"
         assert embedder.api_base_url == "https://custom-api-base-url.com"
         assert embedder.truncate == "START"
@@ -59,7 +59,7 @@ def test_to_dict(self):
         assert component_dict == {
             "type": "cohere_haystack.embedders.text_embedder.CohereTextEmbedder",
             "init_parameters": {
-                "model_name": "embed-english-v2.0",
+                "model": "embed-english-v2.0",
                 "input_type": "search_query",
                 "api_base_url": COHERE_API_URL,
                 "truncate": "END",
@@ -75,7 +75,7 @@ def test_to_dict_with_custom_init_parameters(self):
         """
         embedder_component = CohereTextEmbedder(
             api_key="test-api-key",
-            model_name="embed-multilingual-v2.0",
+            model="embed-multilingual-v2.0",
             input_type="classification",
             api_base_url="https://custom-api-base-url.com",
             truncate="START",
@@ -87,7 +87,7 @@ def test_to_dict_with_custom_init_parameters(self):
         assert component_dict == {
             "type": "cohere_haystack.embedders.text_embedder.CohereTextEmbedder",
             "init_parameters": {
-                "model_name": "embed-multilingual-v2.0",
+                "model": "embed-multilingual-v2.0",
                 "input_type": "classification",
                 "api_base_url": "https://custom-api-base-url.com",
                 "truncate": "START",