Use airflow.models.dag.DAG in Amazon system tests (apache#34613)

asapekia · Sep 26, 2023 · 3b30b8f · 3b30b8f
1 parent 11cba87
commit 3b30b8f
Show file tree

Hide file tree

Showing 47 changed files with 59 additions and 54 deletions.
diff --git a/tests/system/providers/amazon/aws/example_appflow.py b/tests/system/providers/amazon/aws/example_appflow.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.operators.bash import BashOperator
 from airflow.providers.amazon.aws.operators.appflow import (
     AppflowRecordsShortCircuitOperator,

diff --git a/tests/system/providers/amazon/aws/example_appflow_run.py b/tests/system/providers/amazon/aws/example_appflow_run.py
@@ -21,9 +21,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.appflow import (
     AppflowRunOperator,
 )

diff --git a/tests/system/providers/amazon/aws/example_athena.py b/tests/system/providers/amazon/aws/example_athena.py
@@ -20,9 +20,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.operators.athena import AthenaOperator
 from airflow.providers.amazon.aws.operators.s3 import (

diff --git a/tests/system/providers/amazon/aws/example_azure_blob_to_s3.py b/tests/system/providers/amazon/aws/example_azure_blob_to_s3.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.azure_blob_to_s3 import AzureBlobStorageToS3Operator
 from airflow.utils.trigger_rule import TriggerRule

diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py
@@ -21,9 +21,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.batch import BatchCreateComputeEnvironmentOperator, BatchOperator
 from airflow.providers.amazon.aws.sensors.batch import (
     BatchComputeEnvironmentSensor,

diff --git a/tests/system/providers/amazon/aws/example_cloudformation.py b/tests/system/providers/amazon/aws/example_cloudformation.py
@@ -19,8 +19,8 @@
 import json
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.cloud_formation import (
     CloudFormationCreateStackOperator,
     CloudFormationDeleteStackOperator,

diff --git a/tests/system/providers/amazon/aws/example_datasync.py b/tests/system/providers/amazon/aws/example_datasync.py
@@ -20,9 +20,9 @@
 
 import boto3
 
-from airflow import models
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.datasync import DataSyncOperator
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.utils.trigger_rule import TriggerRule
@@ -116,7 +116,7 @@ def delete_locations(locations):
         )
 
 
-with models.DAG(
+with DAG(
     DAG_ID,
     schedule="@once",
     start_date=datetime(2021, 1, 1),

diff --git a/tests/system/providers/amazon/aws/example_dms.py b/tests/system/providers/amazon/aws/example_dms.py
@@ -29,9 +29,9 @@
 import boto3
 from sqlalchemy import Column, MetaData, String, Table, create_engine
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.dms import (
     DmsCreateTaskOperator,
     DmsDeleteTaskOperator,

diff --git a/tests/system/providers/amazon/aws/example_dynamodb.py b/tests/system/providers/amazon/aws/example_dynamodb.py
@@ -20,9 +20,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.sensors.dynamodb import DynamoDBValueSensor
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder

diff --git a/tests/system/providers/amazon/aws/example_ec2.py b/tests/system/providers/amazon/aws/example_ec2.py
@@ -21,9 +21,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.ec2 import (
     EC2CreateInstanceOperator,
     EC2StartInstanceOperator,

diff --git a/tests/system/providers/amazon/aws/example_ecs.py b/tests/system/providers/amazon/aws/example_ecs.py
@@ -20,9 +20,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates
 from airflow.providers.amazon.aws.operators.ecs import (
     EcsCreateClusterOperator,

diff --git a/tests/system/providers/amazon/aws/example_ecs_fargate.py b/tests/system/providers/amazon/aws/example_ecs_fargate.py
@@ -20,9 +20,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.ecs import EcsTaskStates
 from airflow.providers.amazon.aws.operators.ecs import EcsRunTaskOperator
 from airflow.providers.amazon.aws.sensors.ecs import EcsTaskStateSensor

diff --git a/tests/system/providers/amazon/aws/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py
@@ -23,9 +23,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.ssm import SsmHook
 from airflow.providers.amazon.aws.operators.emr import (
     EmrAddStepsOperator,

diff --git a/tests/system/providers/amazon/aws/example_emr_eks.py b/tests/system/providers/amazon/aws/example_emr_eks.py
@@ -22,9 +22,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates
 from airflow.providers.amazon.aws.operators.eks import EksCreateClusterOperator, EksDeleteClusterOperator
 from airflow.providers.amazon.aws.operators.emr import EmrContainerOperator, EmrEksCreateClusterOperator

diff --git a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py
@@ -20,8 +20,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.emr import (
     EmrStartNotebookExecutionOperator,
     EmrStopNotebookExecutionOperator,

diff --git a/tests/system/providers/amazon/aws/example_eventbridge.py b/tests/system/providers/amazon/aws/example_eventbridge.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.eventbridge import (
     EventBridgeDisableRuleOperator,
     EventBridgeEnableRuleOperator,

diff --git a/tests/system/providers/amazon/aws/example_ftp_to_s3.py b/tests/system/providers/amazon/aws/example_ftp_to_s3.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import models
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.ftp_to_s3 import FTPToS3Operator
 from airflow.utils.trigger_rule import TriggerRule
@@ -29,7 +29,7 @@
 
 DAG_ID = "example_ftp_to_s3"
 
-with models.DAG(
+with DAG(
     DAG_ID,
     schedule="@once",
     start_date=datetime(2021, 1, 1),

diff --git a/tests/system/providers/amazon/aws/example_gcs_to_s3.py b/tests/system/providers/amazon/aws/example_gcs_to_s3.py
@@ -18,9 +18,9 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import (
     S3CreateBucketOperator,
     S3DeleteBucketOperator,

diff --git a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py
@@ -20,8 +20,8 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.operators.python import task
 from airflow.providers.amazon.aws.operators.glacier import (
     GlacierCreateJobOperator,

diff --git a/tests/system/providers/amazon/aws/example_glue.py b/tests/system/providers/amazon/aws/example_glue.py
@@ -21,9 +21,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.glue import GlueJobOperator
 from airflow.providers.amazon.aws.operators.glue_crawler import GlueCrawlerOperator
 from airflow.providers.amazon.aws.operators.s3 import (

diff --git a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py
@@ -23,8 +23,8 @@
 import os
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.google_api_to_s3 import GoogleApiToS3Operator
 from airflow.utils.trigger_rule import TriggerRule

diff --git a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py
@@ -52,10 +52,11 @@
 
 import boto3
 
-from airflow import DAG, settings
+from airflow import settings
 from airflow.decorators import task
 from airflow.models import Connection
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.google_api_to_s3 import GoogleApiToS3Operator
 from airflow.utils.trigger_rule import TriggerRule

diff --git a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py
@@ -23,10 +23,10 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models import Connection
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
 from airflow.providers.amazon.aws.transfers.hive_to_dynamodb import HiveToDynamoDBOperator
 from airflow.utils import db

diff --git a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py
@@ -22,8 +22,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.imap_attachment_to_s3 import ImapAttachmentToS3Operator
 from airflow.utils.trigger_rule import TriggerRule

diff --git a/tests/system/providers/amazon/aws/example_lambda.py b/tests/system/providers/amazon/aws/example_lambda.py
@@ -23,9 +23,9 @@
 
 import boto3
 
-from airflow import models
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.lambda_function import (
     LambdaCreateFunctionOperator,
     LambdaInvokeFunctionOperator,
@@ -66,7 +66,7 @@ def delete_lambda(function_name: str):
     )
 
 
-with models.DAG(
+with DAG(
     DAG_ID,
     schedule="@once",
     start_date=datetime(2021, 1, 1),

diff --git a/tests/system/providers/amazon/aws/example_local_to_s3.py b/tests/system/providers/amazon/aws/example_local_to_s3.py
@@ -19,9 +19,9 @@
 import os
 from datetime import datetime
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.local_to_s3 import LocalFilesystemToS3Operator
 from airflow.utils.trigger_rule import TriggerRule

diff --git a/tests/system/providers/amazon/aws/example_mongo_to_s3.py b/tests/system/providers/amazon/aws/example_mongo_to_s3.py
@@ -16,8 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from airflow import models
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator
 from airflow.providers.amazon.aws.transfers.mongo_to_s3 import MongoToS3Operator
 from airflow.utils.dates import datetime
@@ -34,7 +34,7 @@
     SystemTestContextBuilder().add_variable(MONGO_DATABASE_KEY).add_variable(MONGO_COLLECTION_KEY).build()
 )
 
-with models.DAG(
+with DAG(
     DAG_ID,
     schedule="@once",
     start_date=datetime(2021, 1, 1),

diff --git a/tests/system/providers/amazon/aws/example_quicksight.py b/tests/system/providers/amazon/aws/example_quicksight.py
@@ -22,9 +22,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.quicksight import QuickSightCreateIngestionOperator
 from airflow.providers.amazon.aws.operators.s3 import (
     S3CreateBucketOperator,

diff --git a/tests/system/providers/amazon/aws/example_rds_event.py b/tests/system/providers/amazon/aws/example_rds_event.py
@@ -21,9 +21,9 @@
 
 import boto3
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.rds import (
     RdsCreateDbInstanceOperator,
     RdsCreateEventSubscriptionOperator,

diff --git a/tests/system/providers/amazon/aws/example_rds_export.py b/tests/system/providers/amazon/aws/example_rds_export.py
@@ -18,9 +18,9 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.operators.rds import (
     RdsCancelExportTaskOperator,

diff --git a/tests/system/providers/amazon/aws/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.rds import (
     RdsCreateDbInstanceOperator,
     RdsDeleteDbInstanceOperator,

diff --git a/tests/system/providers/amazon/aws/example_rds_snapshot.py b/tests/system/providers/amazon/aws/example_rds_snapshot.py
@@ -18,8 +18,8 @@
 
 from datetime import datetime
 
-from airflow import DAG
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.rds import (
     RdsCopyDbSnapshotOperator,
     RdsCreateDbInstanceOperator,

diff --git a/tests/system/providers/amazon/aws/example_redshift.py b/tests/system/providers/amazon/aws/example_redshift.py
@@ -22,10 +22,11 @@
 
 import boto3
 
-from airflow import DAG, settings
+from airflow import settings
 from airflow.decorators import task
 from airflow.models import Connection
 from airflow.models.baseoperator import chain
+from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.operators.redshift_cluster import (
     RedshiftCreateClusterOperator,