diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index 77fb533258..47b46ed209 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -1577,7 +1577,9 @@ def raw_predict( ) self.raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:rawPredict" - return self.authorized_session.post(self.raw_predict_request_url, body, headers) + return self.authorized_session.post( + url=self.raw_predict_request_url, data=body, headers=headers + ) def explain( self, diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index a528f46981..1c9e7dd4a0 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -61,6 +61,10 @@ from google.protobuf import field_mask_pb2, timestamp_pb2 +from test_endpoints import ( # noqa: F401 + create_endpoint_mock, +) + _TEST_PROJECT = "test-project" _TEST_PROJECT_2 = "test-project-2" _TEST_LOCATION = "us-central1" @@ -2782,5 +2786,7 @@ def test_raw_predict(self, raw_predict_mock): test_endpoint = models.Endpoint(_TEST_ID) test_endpoint.raw_predict(_TEST_RAW_PREDICT_DATA, _TEST_RAW_PREDICT_HEADER) raw_predict_mock.assert_called_once_with( - _TEST_RAW_PREDICT_URL, _TEST_RAW_PREDICT_DATA, _TEST_RAW_PREDICT_HEADER + url=_TEST_RAW_PREDICT_URL, + data=_TEST_RAW_PREDICT_DATA, + headers=_TEST_RAW_PREDICT_HEADER, )