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 <kaxiln...@gmail.com>
Authored: Tue Mar 6 10:22:49 2018 +0100
Committer: Fokko Driesprong <fokkodriespr...@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
+        )

Reply via email to