test: fix sagamaker test cases

astronomer · Jan 25, 2024 · bb61993 · bb61993
1 parent a15b96e
commit bb61993
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/amazon/aws/operators/test_sagemaker.py b/tests/amazon/aws/operators/test_sagemaker.py
@@ -4,6 +4,10 @@
 from airflow.exceptions import AirflowException, TaskDeferred
 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 (
+    SageMakerProcessingOperator,
+    SageMakerTransformOperator,
+)
 from airflow.utils.timezone import datetime
 
 from astronomer.providers.amazon.aws.operators.sagemaker import (
@@ -142,7 +146,7 @@ def test_init(self):
             check_interval=self.CHECK_INTERVAL,
             max_ingestion_time=self.MAX_INGESTION_TIME,
         )
-        assert isinstance(task, SageMakerProcessingOperatorAsync)
+        assert isinstance(task, SageMakerProcessingOperator)
         assert task.deferrable is True
 
 
@@ -159,7 +163,7 @@ def test_init(self):
             check_interval=self.CHECK_INTERVAL,
             max_ingestion_time=self.MAX_INGESTION_TIME,
         )
-        assert isinstance(task, SageMakerProcessingOperatorAsync)
+        assert isinstance(task, SageMakerTransformOperator)
         assert task.deferrable is True