You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by fo...@apache.org on 2018/03/06 09:22:56 UTC
incubator-airflow git commit: [AIRFLOW-2177] Add mock test for GCS
Download op
Repository: incubator-airflow
Updated Branches:
refs/heads/master 976fd1245 -> 587a14b26
[AIRFLOW-2177] Add mock test for GCS Download op
- Added mock test
`GoogleCloudStorageDownloadOperatorTest` for
`GoogleCloudStorageDownloadOperator`
- Minor change in example for
`GoogleCloudStorageCopyOperator`
Closes #3097 from kaxil/gcs-tests
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/587a14b2
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/587a14b2
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/587a14b2
Branch: refs/heads/master
Commit: 587a14b2658a5929b151cd1ce9ed74ab67b250c6
Parents: 976fd12
Author: Kaxil Naik <ka...@gmail.com>
Authored: Tue Mar 6 10:22:49 2018 +0100
Committer: Fokko Driesprong <fo...@godatadriven.com>
Committed: Tue Mar 6 10:22:49 2018 +0100
----------------------------------------------------------------------
airflow/contrib/operators/gcs_copy_operator.py | 6 +--
.../contrib/operators/gcs_download_operator.py | 13 ++++--
.../operators/test_gcs_download_operator.py | 46 ++++++++++++++++++++
3 files changed, 58 insertions(+), 7 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/587a14b2/airflow/contrib/operators/gcs_copy_operator.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/operators/gcs_copy_operator.py b/airflow/contrib/operators/gcs_copy_operator.py
index e679ccb..4922708 100644
--- a/airflow/contrib/operators/gcs_copy_operator.py
+++ b/airflow/contrib/operators/gcs_copy_operator.py
@@ -43,14 +43,14 @@ class GoogleCloudStorageCopyOperator(BaseOperator):
:type delegate_to: string
**Example**:
- The following Operator would move all the CSV files from `sales/sales-2017` folder in
+ The following Operator would copy all the CSV files from `sales/sales-2017` folder in
`data` bucket to `sales` folder in `archive` bucket. ::
move_file = GoogleCloudStorageCopyOperator(
- task_id='move_file',
+ task_id='copy_files',
source_bucket='data',
source_object='sales/sales-2017/',
- source_files_delimiter='.csv'
+ source_files_delimiter='.csv',
destination_bucket='archive',
destination_directory='sales',
google_cloud_storage_conn_id='airflow-service-account'
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/587a14b2/airflow/contrib/operators/gcs_download_operator.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/operators/gcs_download_operator.py b/airflow/contrib/operators/gcs_download_operator.py
index b1df63a..75de608 100644
--- a/airflow/contrib/operators/gcs_download_operator.py
+++ b/airflow/contrib/operators/gcs_download_operator.py
@@ -66,10 +66,15 @@ class GoogleCloudStorageDownloadOperator(BaseOperator):
self.delegate_to = delegate_to
def execute(self, context):
- self.log.info('Executing download: %s, %s, %s', self.bucket, self.object, self.filename)
- hook = GoogleCloudStorageHook(google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
- delegate_to=self.delegate_to)
- file_bytes = hook.download(self.bucket, self.object, self.filename)
+ self.log.info('Executing download: %s, %s, %s', self.bucket,
+ self.object, self.filename)
+ hook = GoogleCloudStorageHook(
+ google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
+ delegate_to=self.delegate_to
+ )
+ file_bytes = hook.download(bucket=self.bucket,
+ object=self.object,
+ filename=self.filename)
if self.store_to_xcom_key:
if sys.getsizeof(file_bytes) < 48000:
context['ti'].xcom_push(key=self.store_to_xcom_key, value=file_bytes)
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/587a14b2/tests/contrib/operators/test_gcs_download_operator.py
----------------------------------------------------------------------
diff --git a/tests/contrib/operators/test_gcs_download_operator.py b/tests/contrib/operators/test_gcs_download_operator.py
new file mode 100644
index 0000000..c234f76
--- /dev/null
+++ b/tests/contrib/operators/test_gcs_download_operator.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from airflow.contrib.operators.gcs_download_operator \
+ import GoogleCloudStorageDownloadOperator
+
+try:
+ from unittest import mock
+except ImportError:
+ try:
+ import mock
+ except ImportError:
+ mock = None
+
+TASK_ID = 'test-gcs-download-operator'
+TEST_BUCKET = 'test-bucket'
+TEST_OBJECT = 'dir1/test-object'
+LOCAL_FILE_PATH = '/home/airflow/gcp/test-object'
+
+
+class GoogleCloudStorageDownloadOperatorTest(unittest.TestCase):
+
+ @mock.patch('airflow.contrib.operators.gcs_download_operator.GoogleCloudStorageHook')
+ def test_execute(self, mock_hook):
+ operator = GoogleCloudStorageDownloadOperator(task_id=TASK_ID,
+ bucket=TEST_BUCKET,
+ object=TEST_OBJECT,
+ filename=LOCAL_FILE_PATH)
+
+ operator.execute(None)
+ mock_hook.return_value.download.assert_called_once_with(
+ bucket=TEST_BUCKET, object=TEST_OBJECT, filename=LOCAL_FILE_PATH
+ )