Fix unit tests

pinecone-io · Aug 22, 2024 · b37d521 · b37d521
1 parent 24fa01f
commit b37d521
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/tests/unit/test_index_initialization.py b/tests/unit/test_index_initialization.py
@@ -5,18 +5,21 @@
 
 class TestIndexClientInitialization:
     @pytest.mark.parametrize("additional_headers", [None, {}])
-    def test_no_additional_headers_leaves_useragent_only(self, additional_headers):
+    def test_no_additional_headers_leaves_useragent_and_version_only(self, additional_headers):
         pc = Pinecone(api_key="YOUR_API_KEY")
         index = pc.Index(host="myhost", additional_headers=additional_headers)
-        assert len(index._vector_api.api_client.default_headers) == 1
+        assert len(index._vector_api.api_client.default_headers) == 2
         assert "User-Agent" in index._vector_api.api_client.default_headers
         assert "python-client-" in index._vector_api.api_client.default_headers["User-Agent"]
+        assert "X-Pinecone-API-Version" in index._vector_api.api_client.default_headers
 
     def test_additional_headers_one_additional(self):
         pc = Pinecone(api_key="YOUR_API_KEY")
         index = pc.Index(host="myhost", additional_headers={"test-header": "test-header-value"})
         assert "test-header" in index._vector_api.api_client.default_headers
-        assert len(index._vector_api.api_client.default_headers) == 2
+        assert "User-Agent" in index._vector_api.api_client.default_headers
+        assert "X-Pinecone-API-Version" in index._vector_api.api_client.default_headers
+        assert len(index._vector_api.api_client.default_headers) == 3
 
     def test_multiple_additional_headers(self):
         pc = Pinecone(api_key="YOUR_API_KEY")
@@ -25,16 +28,17 @@ def test_multiple_additional_headers(self):
         )
         assert "test-header" in index._vector_api.api_client.default_headers
         assert "test-header2" in index._vector_api.api_client.default_headers
-        assert len(index._vector_api.api_client.default_headers) == 3
+        assert len(index._vector_api.api_client.default_headers) == 4
 
     def test_overwrite_useragent(self):
         # This doesn't seem like a common use case, but we may want to allow this
         # when embedding the client in other pinecone tools such as canopy.
         pc = Pinecone(api_key="YOUR_API_KEY")
         index = pc.Index(host="myhost", additional_headers={"User-Agent": "test-user-agent"})
-        assert len(index._vector_api.api_client.default_headers) == 1
+        assert len(index._vector_api.api_client.default_headers) == 2
         assert "User-Agent" in index._vector_api.api_client.default_headers
         assert index._vector_api.api_client.default_headers["User-Agent"] == "test-user-agent"
+        assert "X-Pinecone-API-Version" in index._vector_api.api_client.default_headers
 
     def test_set_source_tag(self):
         pc = Pinecone(api_key="123-456-789", source_tag="test_source_tag")