@@ -84,6 +84,44 @@ def test_execute_copy_single_file(self, sftp_hook, gcs_hook):
84
84
object_name = DESTINATION_PATH_FILE ,
85
85
filename = mock .ANY ,
86
86
mime_type = DEFAULT_MIME_TYPE ,
87
+ gzip = False ,
88
+ )
89
+
90
+ sftp_hook .return_value .delete_file .assert_not_called ()
91
+
92
+ @mock .patch ("airflow.providers.google.cloud.transfers.sftp_to_gcs.GCSHook" )
93
+ @mock .patch ("airflow.providers.google.cloud.transfers.sftp_to_gcs.SFTPHook" )
94
+ def test_execute_copy_single_file_with_compression (self , sftp_hook , gcs_hook ):
95
+ task = SFTPToGCSOperator (
96
+ task_id = TASK_ID ,
97
+ source_path = SOURCE_OBJECT_NO_WILDCARD ,
98
+ destination_bucket = TEST_BUCKET ,
99
+ destination_path = DESTINATION_PATH_FILE ,
100
+ move_object = False ,
101
+ gcp_conn_id = GCP_CONN_ID ,
102
+ sftp_conn_id = SFTP_CONN_ID ,
103
+ delegate_to = DELEGATE_TO ,
104
+ impersonation_chain = IMPERSONATION_CHAIN ,
105
+ gzip = True ,
106
+ )
107
+ task .execute (None )
108
+ gcs_hook .assert_called_once_with (
109
+ gcp_conn_id = GCP_CONN_ID ,
110
+ delegate_to = DELEGATE_TO ,
111
+ impersonation_chain = IMPERSONATION_CHAIN ,
112
+ )
113
+ sftp_hook .assert_called_once_with (SFTP_CONN_ID )
114
+
115
+ sftp_hook .return_value .retrieve_file .assert_called_once_with (
116
+ os .path .join (SOURCE_OBJECT_NO_WILDCARD ), mock .ANY
117
+ )
118
+
119
+ gcs_hook .return_value .upload .assert_called_once_with (
120
+ bucket_name = TEST_BUCKET ,
121
+ object_name = DESTINATION_PATH_FILE ,
122
+ filename = mock .ANY ,
123
+ mime_type = DEFAULT_MIME_TYPE ,
124
+ gzip = True ,
87
125
)
88
126
89
127
sftp_hook .return_value .delete_file .assert_not_called ()
@@ -119,6 +157,7 @@ def test_execute_move_single_file(self, sftp_hook, gcs_hook):
119
157
object_name = DESTINATION_PATH_FILE ,
120
158
filename = mock .ANY ,
121
159
mime_type = DEFAULT_MIME_TYPE ,
160
+ gzip = False ,
122
161
)
123
162
124
163
sftp_hook .return_value .delete_file .assert_called_once_with (SOURCE_OBJECT_NO_WILDCARD )
@@ -162,12 +201,14 @@ def test_execute_copy_with_wildcard(self, sftp_hook, gcs_hook):
162
201
object_name = "destination_dir/test_object3.json" ,
163
202
mime_type = DEFAULT_MIME_TYPE ,
164
203
filename = mock .ANY ,
204
+ gzip = False ,
165
205
),
166
206
mock .call (
167
207
bucket_name = TEST_BUCKET ,
168
208
object_name = "destination_dir/sub_dir/test_object3.json" ,
169
209
mime_type = DEFAULT_MIME_TYPE ,
170
210
filename = mock .ANY ,
211
+ gzip = False ,
171
212
),
172
213
]
173
214
)
0 commit comments