Fix SageMakerEndpointConfigOperator's return value (#26541)

* Fix SageMakerEndpointConfigOperator's JSON return value * Fix the unit tests for the changed operator
apache · Sep 22, 2022 · 0c7b4cb · 0c7b4cb
1 parent 53c0122
commit 0c7b4cb
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/airflow/providers/amazon/aws/operators/sagemaker.py b/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -237,7 +237,11 @@ def execute(self, context: Context) -> dict:
         if response['ResponseMetadata']['HTTPStatusCode'] != 200:
             raise AirflowException(f'Sagemaker endpoint config creation failed: {response}')
         else:
-            return {'EndpointConfig': self.hook.describe_endpoint_config(self.config['EndpointConfigName'])}
+            return {
+                'EndpointConfig': serialize(
+                    self.hook.describe_endpoint_config(self.config['EndpointConfigName'])
+                )
+            }
 
 
 class SageMakerEndpointOperator(SageMakerBaseOperator):

diff --git a/tests/providers/amazon/aws/operators/test_sagemaker_endpoint_config.py b/tests/providers/amazon/aws/operators/test_sagemaker_endpoint_config.py
@@ -24,6 +24,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
+from airflow.providers.amazon.aws.operators import sagemaker
 from airflow.providers.amazon.aws.operators.sagemaker import SageMakerEndpointConfigOperator
 
 CREATE_ENDPOINT_CONFIG_PARAMS: dict = {
@@ -50,7 +51,8 @@ def setUp(self):
 
     @mock.patch.object(SageMakerHook, 'get_conn')
     @mock.patch.object(SageMakerHook, 'create_endpoint_config')
-    def test_integer_fields(self, mock_model, mock_client):
+    @mock.patch.object(sagemaker, 'serialize', return_value="")
+    def test_integer_fields(self, serialize, mock_model, mock_client):
         mock_model.return_value = {
             'EndpointConfigArn': 'test_arn',
             'ResponseMetadata': {'HTTPStatusCode': 200},
@@ -62,7 +64,8 @@ def test_integer_fields(self, mock_model, mock_client):
 
     @mock.patch.object(SageMakerHook, 'get_conn')
     @mock.patch.object(SageMakerHook, 'create_endpoint_config')
-    def test_execute(self, mock_model, mock_client):
+    @mock.patch.object(sagemaker, 'serialize', return_value="")
+    def test_execute(self, serialize, mock_model, mock_client):
         mock_model.return_value = {
             'EndpointConfigArn': 'test_arn',
             'ResponseMetadata': {'HTTPStatusCode': 200},