Repository: apache/airflow

Issue:
Add unittests for GCSTaskHandler
**Description**

The `GCSTaskHandler` from `airflow.utils.log.gcs_task_handler` does not have a unit test. 

**Use case / motivation**

The test is necessary so that we can be able to move the `GCSTaskHandler` to its providers package(`google`)

**Related Issues**
#9386 



Discussion:
Hey @ephraimbuddy are you working on it or is this issue up for taking?
You can work on it, I'm not currently working on it.
Great, can you assign me this task please. Thanks :smile: 

Resolving patch:
diff --git a/airflow/providers/google/cloud/log/gcs_task_handler.py b/airflow/providers/google/cloud/log/gcs_task_handler.py
--- a/airflow/providers/google/cloud/log/gcs_task_handler.py
+++ b/airflow/providers/google/cloud/log/gcs_task_handler.py
@@ -133,7 +133,7 @@ def gcs_read(self, remote_log_location):
         bkt, blob = self.parse_gcs_url(remote_log_location)
         return self.hook.download(bkt, blob).decode('utf-8')
 
-    def gcs_write(self, log, remote_log_location, append=True):
+    def gcs_write(self, log, remote_log_location):
         """
         Writes the log to the remote_log_location. Fails silently if no hook
         was created.
@@ -142,28 +142,17 @@ def gcs_write(self, log, remote_log_location, append=True):
         :type log: str
         :param remote_log_location: the log's location in remote storage
         :type remote_log_location: str (path)
-        :param append: if False, any existing log file is overwritten. If True,
-            the new log is appended to any existing logs.
-        :type append: bool
         """
-        if append:
-            try:
-                old_log = self.gcs_read(remote_log_location)
-                log = '\n'.join([old_log, log]) if old_log else log
-            except Exception as e:  # pylint: disable=broad-except
-                if not hasattr(e, 'resp') or e.resp.get('status') != '404':  # pylint: disable=no-member
-                    log = '*** Previous log discarded: {}\n\n'.format(str(e)) + log
+        try:
+            old_log = self.gcs_read(remote_log_location)
+            log = '\n'.join([old_log, log]) if old_log else log
+        except Exception as e:  # pylint: disable=broad-except
+            if not hasattr(e, 'resp') or e.resp.get('status') != '404':  # pylint: disable=no-member
+                log = '*** Previous log discarded: {}\n\n'.format(str(e)) + log
 
         try:
             bkt, blob = self.parse_gcs_url(remote_log_location)
-            from tempfile import NamedTemporaryFile
-            with NamedTemporaryFile(mode='w+') as tmpfile:
-                tmpfile.write(log)
-                # Force the file to be flushed, since we're doing the
-                # upload from within the file context (it hasn't been
-                # closed).
-                tmpfile.flush()
-                self.hook.upload(bkt, blob, tmpfile.name)
+            self.hook.upload(bkt, blob, data=log)
         except Exception as e:  # pylint: disable=broad-except
             self.log.error('Could not write logs to %s: %s', remote_log_location, e)
 


Test patch:
diff --git a/tests/providers/google/cloud/log/test_gcs_task_handler.py b/tests/providers/google/cloud/log/test_gcs_task_handler.py
new file mode 100644
--- /dev/null
+++ b/tests/providers/google/cloud/log/test_gcs_task_handler.py
@@ -0,0 +1,162 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 logging
+import shutil
+import tempfile
+import unittest
+from datetime import datetime
+from unittest import mock
+
+from airflow.models import TaskInstance
+from airflow.models.dag import DAG
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.cloud.log.gcs_task_handler import GCSTaskHandler
+from airflow.utils.state import State
+from tests.test_utils.config import conf_vars
+from tests.test_utils.db import clear_db_runs
+
+
+class TestGCSTaskHandler(unittest.TestCase):
+    def setUp(self) -> None:
+        date = datetime(2020, 1, 1)
+        self.gcs_log_folder = "test"
+        self.logger = logging.getLogger("logger")
+        self.dag = DAG("dag_for_testing_task_handler", start_date=date)
+        task = DummyOperator(task_id="task_for_testing_gcs_task_handler")
+        self.ti = TaskInstance(task=task, execution_date=date)
+        self.ti.try_number = 1
+        self.ti.state = State.RUNNING
+        self.remote_log_base = "gs://bucket/remote/log/location"
+        self.remote_log_location = "gs://my-bucket/path/to/1.log"
+        self.local_log_location = tempfile.mkdtemp()
+        self.filename_template = "{try_number}.log"
+        self.addCleanup(self.dag.clear)
+        self.gcs_task_handler = GCSTaskHandler(
+            self.local_log_location, self.remote_log_base, self.filename_template
+        )
+
+    def tearDown(self) -> None:
+        clear_db_runs()
+        shutil.rmtree(self.local_log_location, ignore_errors=True)
+
+    def test_hook(self):
+        self.assertIsInstance(self.gcs_task_handler.hook, GCSHook)
+
+    @conf_vars({("logging", "remote_log_conn_id"): "gcs_default"})
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_hook_raises(self, mock_hook):
+        mock_hook.side_effect = Exception("Failed to connect")
+
+        with self.assertLogs(self.gcs_task_handler.log) as cm:
+            self.gcs_task_handler.hook
+
+        self.assertEqual(
+            cm.output,
+            ['ERROR:airflow.providers.google.cloud.log.gcs_task_handler.GCSTaskHandler:Could '
+             'not create a GoogleCloudStorageHook with connection id "gcs_default". Failed '
+             'to connect\n'
+             '\n'
+             'Please make sure that airflow[gcp] is installed and the GCS connection '
+             'exists.']
+        )
+
+    @conf_vars({("logging", "remote_log_conn_id"): "gcs_default"})
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_should_read_logs_from_remote(self, mock_hook):
+        mock_hook.return_value.download.return_value = b"CONTENT"
+
+        logs, metadata = self.gcs_task_handler._read(self.ti, self.ti.try_number)
+
+        mock_hook.return_value.download.assert_called_once_with('bucket', 'remote/log/location/1.log')
+        self.assertEqual(
+            '*** Reading remote log from gs://bucket/remote/log/location/1.log.\nCONTENT\n', logs)
+        self.assertEqual({'end_of_log': True}, metadata)
+
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_should_read_from_local(self, mock_hook):
+        mock_hook.return_value.download.side_effect = Exception("Failed to connect")
+
+        self.gcs_task_handler.set_context(self.ti)
+        return_val = self.gcs_task_handler._read(self.ti, self.ti.try_number)
+
+        self.assertEqual(len(return_val), 2)
+        self.assertEqual(
+            return_val[0],
+            "*** Unable to read remote log from gs://bucket/remote/log/location/1.log\n*** "
+            f"Failed to connect\n\n*** Reading local file: {self.local_log_location}/1.log\n",
+        )
+        self.assertDictEqual(return_val[1], {"end_of_log": True})
+        mock_hook.return_value.download.assert_called_once()
+
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_write_to_remote_on_close(self, mock_hook):
+        mock_hook.return_value.download.return_value = b"CONTENT"
+
+        self.gcs_task_handler.set_context(self.ti)
+        self.gcs_task_handler.emit(logging.LogRecord(
+            name="NAME", level="DEBUG", pathname=None, lineno=None,
+            msg="MESSAGE", args=None, exc_info=None
+        ))
+        self.gcs_task_handler.close()
+
+        mock_hook.return_value.download.assert_called_once_with('bucket', 'remote/log/location/1.log')
+        mock_hook.return_value.upload.assert_called_once_with(
+            'bucket', 'remote/log/location/1.log', data='CONTENT\nMESSAGE\n'
+        )
+        self.assertEqual(self.gcs_task_handler.closed, True)
+
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_failed_write_to_remote_on_close(self, mock_hook):
+        mock_hook.return_value.upload.side_effect = Exception("Failed to connect")
+        mock_hook.return_value.download.return_value = b"Old log"
+
+        self.gcs_task_handler.set_context(self.ti)
+        with self.assertLogs(self.gcs_task_handler.log) as cm:
+            self.gcs_task_handler.close()
+
+        self.assertEqual(
+            cm.output,
+            [
+                'ERROR:airflow.providers.google.cloud.log.gcs_task_handler.GCSTaskHandler:Could '
+                'not write logs to gs://bucket/remote/log/location/1.log: Failed to connect'
+            ]
+        )
+        mock_hook.return_value.download.assert_called_once_with(
+            'bucket', 'remote/log/location/1.log'
+        )
+        mock_hook.return_value.upload.assert_called_once_with(
+            'bucket', 'remote/log/location/1.log', data='Old log\n'
+        )
+
+    @mock.patch("airflow.providers.google.cloud.hooks.gcs.GCSHook")
+    def test_write_to_remote_on_close_failed_read_old_logs(self, mock_hook):
+        mock_hook.return_value.download.side_effect = Exception("Fail to download")
+
+        self.gcs_task_handler.set_context(self.ti)
+        self.gcs_task_handler.emit(logging.LogRecord(
+            name="NAME", level="DEBUG", pathname=None, lineno=None,
+            msg="MESSAGE", args=None, exc_info=None
+        ))
+        self.gcs_task_handler.close()
+
+        mock_hook.return_value.download.assert_called_once_with('bucket', 'remote/log/location/1.log')
+        mock_hook.return_value.upload.assert_called_once_with(
+            'bucket', 'remote/log/location/1.log',
+            data='*** Previous log discarded: Fail to download\n\nMESSAGE\n'
+        )
+        self.assertEqual(self.gcs_task_handler.closed, True)
diff --git a/tests/providers/google/cloud/log/test_gcs_task_handler_system.py b/tests/providers/google/cloud/log/test_gcs_task_handler_system.py
new file mode 100644
--- /dev/null
+++ b/tests/providers/google/cloud/log/test_gcs_task_handler_system.py
@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 importlib
+import random
+import string
+import subprocess
+from unittest import mock
+
+import pytest
+
+from airflow import settings
+from airflow.example_dags import example_complex
+from airflow.models import DagBag, TaskInstance
+from airflow.utils.log.log_reader import TaskLogReader
+from airflow.utils.session import provide_session
+from tests.providers.google.cloud.utils.gcp_authenticator import GCP_GCS_KEY
+from tests.test_utils.config import conf_vars
+from tests.test_utils.db import clear_db_connections, clear_db_runs
+from tests.test_utils.gcp_system_helpers import (
+    GoogleSystemTest, provide_gcp_context, resolve_full_gcp_key_path,
+)
+
+
+@pytest.mark.system("google")
+@pytest.mark.credential_file(GCP_GCS_KEY)
+class TestGCSTaskHandlerSystemTest(GoogleSystemTest):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        unique_suffix = ''.join(random.sample(string.ascii_lowercase, 16))
+        cls.bucket_name = f"airflow-gcs-task-handler-tests-{unique_suffix}"  # type: ignore
+        cls.create_gcs_bucket(cls.bucket_name)  # type: ignore
+        clear_db_connections()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls.delete_gcs_bucket(cls.bucket_name)  # type: ignore
+
+    def setUp(self) -> None:
+        clear_db_runs()
+
+    def tearDown(self) -> None:
+        from airflow.config_templates import airflow_local_settings
+        importlib.reload(airflow_local_settings)
+        settings.configure_logging()
+        clear_db_runs()
+
+    @provide_session
+    def test_should_read_logs(self, session):
+        with mock.patch.dict(
+            'os.environ',
+            AIRFLOW__LOGGING__REMOTE_LOGGING="true",
+            AIRFLOW__LOGGING__REMOTE_BASE_LOG_FOLDER=f"gs://{self.bucket_name}/path/to/logs",
+            AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID="google_cloud_default",
+            AIRFLOW__CORE__LOAD_EXAMPLES="false",
+            AIRFLOW__CORE__DAGS_FOLDER=example_complex.__file__,
+            GOOGLE_APPLICATION_CREDENTIALS=resolve_full_gcp_key_path(GCP_GCS_KEY)
+        ):
+            self.assertEqual(0, subprocess.Popen(
+                ["airflow", "dags", "trigger", "example_complex"]
+            ).wait())
+            self.assertEqual(0, subprocess.Popen(
+                ["airflow", "scheduler", "--num-runs", "1"]
+            ).wait())
+
+        ti = session.query(TaskInstance).filter(TaskInstance.task_id == "create_entry_group").first()
+        dag = DagBag(dag_folder=example_complex.__file__).dags['example_complex']
+        task = dag.task_dict["create_entry_group"]
+        ti.task = task
+        self.assert_remote_logs("INFO - Task exited with return code 0", ti)
+
+    def assert_remote_logs(self, expected_message, ti):
+        with provide_gcp_context(GCP_GCS_KEY), conf_vars({
+            ('logging', 'remote_logging'): 'True',
+            ('logging', 'remote_base_log_folder'): f"gs://{self.bucket_name}/path/to/logs",
+            ('logging', 'remote_log_conn_id'): "google_cloud_default",
+        }):
+            from airflow.config_templates import airflow_local_settings
+            importlib.reload(airflow_local_settings)
+            settings.configure_logging()
+
+            task_log_reader = TaskLogReader()
+            logs = "\n".join(task_log_reader.read_log_stream(ti, try_number=None, metadata={}))
+            self.assertIn(expected_message, logs)
diff --git a/tests/test_project_structure.py b/tests/test_project_structure.py
--- a/tests/test_project_structure.py
+++ b/tests/test_project_structure.py
@@ -28,8 +28,7 @@
 )
 
 MISSING_TEST_FILES = {
-    'tests/providers/google/cloud/log/test_gcs_task_handler.py',
-    'tests/providers/microsoft/azure/sensors/test_azure_cosmos.py'
+    'tests/providers/microsoft/azure/sensors/test_azure_cosmos.py',
 }