Fix integrationId casing, add tests

pinecone/core_ea/openapi/db_data/model/start_import_request.py

@@ -112,7 +112,7 @@ def discriminator():
         return None
 
     attribute_map = {
-        "integration_id": "integration_id",  # noqa: E501
+        "integration_id": "integrationId",  # noqa: E501
         "uri": "uri",  # noqa: E501
         "error_mode": "errorMode",  # noqa: E501
     }

tests/unit/data/test_bulk_import.py

@@ -8,7 +8,7 @@
 from pinecone.core_ea.openapi.shared.api_client import ApiClient
 from pinecone.core_ea.openapi.shared.exceptions import PineconeApiException
 
-from pinecone.data.features.bulk_import import ImportFeatureMixin
+from pinecone.data.features.bulk_import import ImportFeatureMixin, ImportErrorMode
 
 
 def build_api_w_faked_response(mocker, body: str, status: int = 200) -> BaseHTTPResponse:
@@ -18,26 +18,31 @@ def build_api_w_faked_response(mocker, body: str, status: int = 200) -> BaseHTTP
     response.data = body.encode("utf-8")
 
     api_client = ApiClient()
-    mocker.patch.object(api_client.rest_client.pool_manager, "request", return_value=response)
-    return BulkOperationsApi(api_client=api_client)
+    mock_request = mocker.patch.object(api_client.rest_client.pool_manager, "request", return_value=response)
+    return BulkOperationsApi(api_client=api_client), mock_request
 
 
 def build_client_w_faked_response(mocker, body: str, status: int = 200):
-    api_client = build_api_w_faked_response(mocker, body, status)
-    return ImportFeatureMixin(__import_operations_api=api_client, api_key="asdf", host="asdf")
+    api_client, mock_req = build_api_w_faked_response(mocker, body, status)
+    return ImportFeatureMixin(__import_operations_api=api_client, api_key="asdf", host="asdf"), mock_req
 
 
 class TestBulkImportStartImport:
-    def test_start_import(self, mocker):
+    def test_start_import_minimal(self, mocker):
         body = """
         {
             "id": "1"
         }
         """
-        client = build_client_w_faked_response(mocker, body)
+        client, mock_req = build_client_w_faked_response(mocker, body)
 
         with pytest.warns(UserWarning, match="pre-release"):
             my_import = client.start_import("s3://path/to/file.parquet")
+
+            # We made some overrides to the print behavior, so we need to
+            # call it to ensure it doesn't raise an exception
+            print(my_import)
+
             assert my_import.id == "1"
             assert my_import["id"] == "1"
             assert my_import.to_dict() == {"id": "1"}
@@ -49,15 +54,71 @@ def test_start_import_with_kwargs(self, mocker):
             "id": "1"
         }
         """
-        client = build_client_w_faked_response(mocker, body)
+        client, mock_req = build_client_w_faked_response(mocker, body)
 
         with pytest.warns(UserWarning, match="pre-release"):
-            my_import = client.start_import(uri="s3://path/to/file.parquet")
+            my_import = client.start_import(uri="s3://path/to/file.parquet", integration_id="123-456-789")
             assert my_import.id == "1"
             assert my_import["id"] == "1"
             assert my_import.to_dict() == {"id": "1"}
             assert my_import.__class__ == StartImportResponse
 
+            # By default, use continue error mode
+            _, call_kwargs = mock_req.call_args
+            assert (
+                call_kwargs["body"]
+                == '{"uri": "s3://path/to/file.parquet", "integrationId": "123-456-789", "errorMode": {"onError": "continue"}}'
+            )
+
+    @pytest.mark.parametrize(
+        "error_mode_input",
+        [
+            ImportErrorMode.CONTINUE,
+            "Continue",
+            "continue",
+            "cONTINUE",
+        ],
+    )
+    def test_start_import_with_explicit_error_mode(self, mocker, error_mode_input):
+        body = """
+        {
+            "id": "1"
+        }
+        """
+        client, mock_req = build_client_w_faked_response(mocker, body)
+
+        with pytest.warns(UserWarning, match="pre-release"):
+            my_import = client.start_import(uri="s3://path/to/file.parquet", error_mode=error_mode_input)
+            _, call_kwargs = mock_req.call_args
+            assert call_kwargs["body"] == '{"uri": "s3://path/to/file.parquet", "errorMode": {"onError": "continue"}}'
+
+    def test_start_import_with_abort_error_mode(self, mocker):
+        body = """
+        {
+            "id": "1"
+        }
+        """
+        client, mock_req = build_client_w_faked_response(mocker, body)
+
+        with pytest.warns(UserWarning, match="pre-release"):
+            my_import = client.start_import(uri="s3://path/to/file.parquet", error_mode=ImportErrorMode.ABORT)
+            _, call_kwargs = mock_req.call_args
+            assert call_kwargs["body"] == '{"uri": "s3://path/to/file.parquet", "errorMode": {"onError": "abort"}}'
+
+    def test_start_import_with_unknown_error_mode(self, mocker):
+        body = """
+        {
+            "id": "1"
+        }
+        """
+        client, mock_req = build_client_w_faked_response(mocker, body)
+
+        with pytest.warns(UserWarning, match="pre-release"):
+            with pytest.raises(ValueError) as e:
+                my_import = client.start_import(uri="s3://path/to/file.parquet", error_mode="unknown")
+
+            assert "Invalid error_mode value: unknown" in str(e.value)
+
     def test_start_invalid_uri(self, mocker):
         body = """
         {
@@ -66,7 +127,7 @@ def test_start_invalid_uri(self, mocker):
             "details": []
         }
         """
-        client = build_client_w_faked_response(mocker, body, 400)
+        client, mock_req = build_client_w_faked_response(mocker, body, 400)
 
         with pytest.warns(UserWarning, match="pre-release"):
             with pytest.raises(PineconeApiException) as e:
@@ -77,7 +138,7 @@ def test_start_invalid_uri(self, mocker):
         assert "Bulk import URIs must start with the scheme of a supported storage provider" in str(e.value)
 
     def test_no_arguments(self, mocker):
-        client = build_client_w_faked_response(mocker, "")
+        client, mock_req = build_client_w_faked_response(mocker, "")
 
         with pytest.warns(UserWarning, match="pre-release"):
             with pytest.raises(TypeError) as e:
@@ -102,14 +163,19 @@ def test_describe_import(self, mocker):
             "percent_complete": 43.2
         }
         """
-        client = build_client_w_faked_response(mocker, body)
+        client, mock_req = build_client_w_faked_response(mocker, body)
 
         with pytest.warns(UserWarning, match="pre-release"):
             my_import = client.describe_import(id="1")
+
+            # We made some overrides to the print behavior, so we need to
+            # call it to ensure it doesn't raise an exception
+            print(my_import)
+
             assert my_import.id == "1"
             assert my_import["id"] == "1"
             desc = my_import.to_dict()
-            assert desc['id'] == "1"
-            assert desc['records_imported'] == 1000
-            assert desc['uri'] == "s3://path/to/file.parquet"
-            assert desc['status'] == "InProgress"
+            assert desc["id"] == "1"
+            assert desc["records_imported"] == 1000
+            assert desc["uri"] == "s3://path/to/file.parquet"
+            assert desc["status"] == "InProgress"