From bb619933e71925461df96c2fb9eb315fdcbc7351 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Wed, 24 Jan 2024 14:03:53 +0800 Subject: [PATCH] test: fix sagamaker test cases --- tests/amazon/aws/operators/test_sagemaker.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/amazon/aws/operators/test_sagemaker.py b/tests/amazon/aws/operators/test_sagemaker.py index 97c0ad066..8f25c50e4 100644 --- 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