Skip to content

Commit b4cef6d

Browse files
authored
SFTP to Google Cloud Storage Transfer system tests migration (AIP-47) (#26799)
* SFTP to Google Cloud Storage Transfer system tests migration (AIP-47)
1 parent 618967e commit b4cef6d

File tree

4 files changed

+67
-87
lines changed

4 files changed

+67
-87
lines changed

β€Ždocs/apache-airflow-providers-google/operators/transfer/sftp_to_gcs.rst

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ Copying single files
4646

4747
The following Operator copies a single file.
4848

49-
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_sftp_to_gcs.py
49+
.. exampleinclude:: /../../tests/system/providers/google/cloud/gcs/example_sftp_to_gcs.py
5050
:language: python
5151
:dedent: 4
5252
:start-after: [START howto_operator_sftp_to_gcs_copy_single_file]
@@ -59,7 +59,7 @@ To move the file use the ``move_object`` parameter. Once the file is copied to G
5959
the original file from the SFTP is deleted.
6060
The ``destination_path`` parameter defines the full path of the file in the bucket.
6161

62-
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_sftp_to_gcs.py
62+
.. exampleinclude:: /../../tests/system/providers/google/cloud/gcs/example_sftp_to_gcs.py
6363
:language: python
6464
:dedent: 4
6565
:start-after: [START howto_operator_sftp_to_gcs_move_single_file_destination]
@@ -71,7 +71,7 @@ Copying directory
7171

7272
Use the ``wildcard`` in ``source_path`` parameter to copy the directory.
7373

74-
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_sftp_to_gcs.py
74+
.. exampleinclude:: /../../tests/system/providers/google/cloud/gcs/example_sftp_to_gcs.py
7575
:language: python
7676
:dedent: 4
7777
:start-after: [START howto_operator_sftp_to_gcs_copy_directory]
@@ -87,7 +87,7 @@ e.g. ``tests_sftp_hook_dir/subdir/parent-1.bin`` is copied to ``specific_files/p
8787
and ``tests_sftp_hook_dir/subdir/parent-2.bin`` is copied to ``specific_files/parent-2.bin`` .
8888
``tests_sftp_hook_dir/subdir/parent-3.txt`` is skipped.
8989

90-
.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_sftp_to_gcs.py
90+
.. exampleinclude:: /../../tests/system/providers/google/cloud/gcs/example_sftp_to_gcs.py
9191
:language: python
9292
:dedent: 4
9393
:start-after: [START howto_operator_sftp_to_gcs_move_specific_files]

β€Žtests/providers/google/cloud/transfers/test_sftp_to_gcs_system.py

Lines changed: 0 additions & 71 deletions
This file was deleted.

β€Žairflow/providers/google/cloud/example_dags/example_sftp_to_gcs.py renamed to β€Žtests/system/providers/google/cloud/gcs/example_sftp_to_gcs.py

Lines changed: 63 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,39 +22,61 @@
2222

2323
import os
2424
from datetime import datetime
25+
from pathlib import Path
2526

2627
from airflow import models
28+
from airflow.models.baseoperator import chain
29+
from airflow.operators.bash import BashOperator
30+
from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
2731
from airflow.providers.google.cloud.transfers.sftp_to_gcs import SFTPToGCSOperator
32+
from airflow.utils.trigger_rule import TriggerRule
2833

29-
BUCKET_SRC = os.environ.get("GCP_GCS_BUCKET_1_SRC", "test-sftp-gcs")
34+
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
35+
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
3036

31-
TMP_PATH = "/tmp"
37+
DAG_ID = "example_sftp_to_gcs"
38+
BUCKET_NAME = f"bucket-{DAG_ID}-{ENV_ID}"
39+
40+
TMP_PATH = "tmp"
3241
DIR = "tests_sftp_hook_dir"
3342
SUBDIR = "subdir"
3443

3544
OBJECT_SRC_1 = "parent-1.bin"
3645
OBJECT_SRC_2 = "parent-2.bin"
37-
OBJECT_SRC_3 = "parent-3.txt"
46+
47+
CURRENT_FOLDER = Path(__file__).parent
48+
LOCAL_PATH = str(Path(CURRENT_FOLDER) / "resources")
49+
50+
FILE_LOCAL_PATH = str(Path(LOCAL_PATH) / TMP_PATH / DIR)
51+
FILE_NAME = "tmp.tar.gz"
3852

3953

4054
with models.DAG(
41-
"example_sftp_to_gcs",
55+
DAG_ID,
56+
schedule="@once",
4257
start_date=datetime(2021, 1, 1),
4358
catchup=False,
4459
) as dag:
60+
61+
create_bucket = GCSCreateBucketOperator(task_id="create_bucket", bucket_name=BUCKET_NAME)
62+
63+
unzip_file = BashOperator(
64+
task_id="unzip_data_file", bash_command=f"tar xvf {LOCAL_PATH}/{FILE_NAME} -C {LOCAL_PATH}"
65+
)
66+
4567
# [START howto_operator_sftp_to_gcs_copy_single_file]
4668
copy_file_from_sftp_to_gcs = SFTPToGCSOperator(
4769
task_id="file-copy-sftp-to-gcs",
48-
source_path=os.path.join(TMP_PATH, DIR, OBJECT_SRC_1),
49-
destination_bucket=BUCKET_SRC,
70+
source_path=f"{FILE_LOCAL_PATH}/{OBJECT_SRC_1}",
71+
destination_bucket=BUCKET_NAME,
5072
)
5173
# [END howto_operator_sftp_to_gcs_copy_single_file]
5274

5375
# [START howto_operator_sftp_to_gcs_move_single_file_destination]
5476
move_file_from_sftp_to_gcs_destination = SFTPToGCSOperator(
5577
task_id="file-move-sftp-to-gcs-destination",
56-
source_path=os.path.join(TMP_PATH, DIR, OBJECT_SRC_2),
57-
destination_bucket=BUCKET_SRC,
78+
source_path=f"{FILE_LOCAL_PATH}/{OBJECT_SRC_2}",
79+
destination_bucket=BUCKET_NAME,
5880
destination_path="destination_dir/destination_filename.bin",
5981
move_object=True,
6082
)
@@ -63,17 +85,46 @@
6385
# [START howto_operator_sftp_to_gcs_copy_directory]
6486
copy_directory_from_sftp_to_gcs = SFTPToGCSOperator(
6587
task_id="dir-copy-sftp-to-gcs",
66-
source_path=os.path.join(TMP_PATH, DIR, SUBDIR, "*"),
67-
destination_bucket=BUCKET_SRC,
88+
source_path=f"{FILE_LOCAL_PATH}/{SUBDIR}/*",
89+
destination_bucket=BUCKET_NAME,
6890
)
6991
# [END howto_operator_sftp_to_gcs_copy_directory]
7092

7193
# [START howto_operator_sftp_to_gcs_move_specific_files]
7294
move_specific_files_from_gcs_to_sftp = SFTPToGCSOperator(
7395
task_id="dir-move-specific-files-sftp-to-gcs",
74-
source_path=os.path.join(TMP_PATH, DIR, SUBDIR, "*.bin"),
75-
destination_bucket=BUCKET_SRC,
96+
source_path=f"{FILE_LOCAL_PATH}/{SUBDIR}/*.bin",
97+
destination_bucket=BUCKET_NAME,
7698
destination_path="specific_files/",
7799
move_object=True,
78100
)
79101
# [END howto_operator_sftp_to_gcs_move_specific_files]
102+
103+
delete_bucket = GCSDeleteBucketOperator(
104+
task_id="delete_bucket", bucket_name=BUCKET_NAME, trigger_rule=TriggerRule.ALL_DONE
105+
)
106+
107+
chain(
108+
# TEST SETUP
109+
create_bucket,
110+
unzip_file,
111+
# TEST BODY
112+
copy_file_from_sftp_to_gcs,
113+
move_file_from_sftp_to_gcs_destination,
114+
copy_directory_from_sftp_to_gcs,
115+
move_specific_files_from_gcs_to_sftp,
116+
# TEST TEARDOWN
117+
delete_bucket,
118+
)
119+
120+
from tests.system.utils.watcher import watcher
121+
122+
# This test needs watcher in order to properly mark success/failure
123+
# when "tearDown" task with trigger rule is part of the DAG
124+
list(dag.tasks) >> watcher()
125+
126+
127+
from tests.system.utils import get_test_run # noqa: E402
128+
129+
# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
130+
test_run = get_test_run(dag)
Binary file not shown.

0 commit comments

Comments
 (0)