he-gong-sardine commented on code in PR #36545:
URL: https://github.com/apache/airflow/pull/36545#discussion_r1500592065
##########
tests/providers/google/cloud/transfers/test_sftp_to_gcs.py:
##########
@@ -252,3 +254,60 @@ def test_execute_more_than_one_wildcard_exception(self,
sftp_hook, gcs_hook):
err = ctx.value
assert "Only one wildcard '*' is allowed in source_path parameter" in
str(err)
+
+class TestSFTPToGCSOperatorStream:
+
+ def setup_method(self):
+ # setup @mock.patch
+ patcher_sftp =
mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.SFTPHook")
+ self.mock_sftp_hook = patcher_sftp.start()
+ patcher_gcs =
mock.patch("airflow.providers.google.cloud.transfers.sftp_to_gcs.GCSHook")
+ self.mock_gcs_hook = patcher_gcs.start()
+
+ self.task = SFTPToGCSOperator(
+ task_id="test_task",
+ source_path=SOURCE_OBJECT_NO_WILDCARD,
+ destination_bucket=TEST_BUCKET,
+ destination_path=DESTINATION_PATH_FILE,
+ use_stream=True,
+ sftp_conn_id=SFTP_CONN_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ )
+
+ def teardown_method(self):
+ self.mock_sftp_hook.stop()
+ self.mock_gcs_hook.stop()
+
+ def test_stream_single_object_default_method(self):
+ # Use 'upload_from_file' method by default
+ mock_dest_blob, mock_temp_dest_blob = MagicMock(), MagicMock()
+
self.mock_gcs_hook.return_value.get_conn.return_value.bucket.return_value.blob.side_effect
= [mock_dest_blob, mock_temp_dest_blob]
+ self.task.execute(None)
+ mock_temp_dest_blob.upload_from_file.assert_called()
+
+ def test_stream_single_object_getfo_method(self):
+ mock_dest_blob, mock_temp_dest_blob = MagicMock(), MagicMock()
Review Comment:
thanks. resolved
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
