def test_write_raises(self):
        handler = WasbTaskHandler(
            self.local_log_location, self.wasb_log_folder, self.container_name, self.filename_template, True
        )
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") as mock_hook:
                mock_hook.return_value.load_string.side_effect = AzureHttpError("failed to connect", 404)

                handler.wasb_write('text', self.remote_log_location, append=False)

            mock_error.assert_called_once_with(
                'Could not write logs to %s', 'remote/log/location/1.log', exc_info=True
            )
class TestWasbTaskHandler(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.wasb_log_folder = 'wasb://container/remote/log/location'
        self.remote_log_location = 'remote/log/location/1.log'
        self.local_log_location = 'local/log/location'
        self.container_name = "wasb-container"
        self.filename_template = '{try_number}.log'
        self.wasb_task_handler = WasbTaskHandler(
            base_log_folder=self.local_log_location,
            wasb_log_folder=self.wasb_log_folder,
            wasb_container=self.container_name,
            filename_template=self.filename_template,
            delete_local_copy=True,
        )

        date = datetime(2020, 8, 10)
        self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
        task = DummyOperator(task_id='task_for_testing_file_log_handler',
                             dag=self.dag)
        self.ti = TaskInstance(task=task, execution_date=date)
        self.ti.try_number = 1
        self.ti.state = State.RUNNING
        self.addCleanup(self.dag.clear)

    @conf_vars({('logging', 'remote_log_conn_id'): 'wasb_default'})
    @mock.patch(
        "airflow.providers.microsoft.azure.hooks.wasb.BlobServiceClient")
    def test_hook(self, mock_service):
        assert isinstance(self.wasb_task_handler.hook, WasbHook)

    @conf_vars({('logging', 'remote_log_conn_id'): 'wasb_default'})
    def test_hook_raises(self):
        handler = WasbTaskHandler(self.local_log_location,
                                  self.wasb_log_folder, self.container_name,
                                  self.filename_template, True)
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch(
                    "airflow.providers.microsoft.azure.hooks.wasb.WasbHook"
            ) as mock_hook:
                mock_hook.side_effect = AzureHttpError("failed to connect",
                                                       404)
                # Initialize the hook
                handler.hook

            mock_error.assert_called_once_with(
                'Could not create an WasbHook with connection id "%s".'
                ' Please make sure that apache-airflow[azure] is installed'
                ' and the Wasb connection exists.',
                'wasb_default',
                exc_info=True,
            )

    def test_set_context_raw(self):
        self.ti.raw = True
        self.wasb_task_handler.set_context(self.ti)
        assert not self.wasb_task_handler.upload_on_close

    def test_set_context_not_raw(self):
        self.wasb_task_handler.set_context(self.ti)
        assert self.wasb_task_handler.upload_on_close

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_wasb_log_exists(self, mock_hook):
        instance = mock_hook.return_value
        instance.check_for_blob.return_value = True
        self.wasb_task_handler.wasb_log_exists(self.remote_log_location)
        mock_hook.return_value.check_for_blob.assert_called_once_with(
            self.container_name, self.remote_log_location)

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_wasb_read(self, mock_hook):
        mock_hook.return_value.read_file.return_value = 'Log line'
        assert self.wasb_task_handler.wasb_read(
            self.remote_log_location) == "Log line"
        assert self.wasb_task_handler.read(self.ti) == (
            [[(
                '',
                '*** Reading remote log from wasb://container/remote/log/location/1.log.\n'
                'Log line\n',
            )]],
            [{
                'end_of_log': True
            }],
        )

    def test_wasb_read_raises(self):
        handler = WasbTaskHandler(self.local_log_location,
                                  self.wasb_log_folder, self.container_name,
                                  self.filename_template, True)
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch(
                    "airflow.providers.microsoft.azure.hooks.wasb.WasbHook"
            ) as mock_hook:
                mock_hook.return_value.read_file.side_effect = AzureHttpError(
                    "failed to connect", 404)

                handler.wasb_read(self.remote_log_location, return_error=True)
            mock_error.assert_called_once_with(
                'Could not read logs from remote/log/location/1.log',
                exc_info=True,
            )

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    @mock.patch.object(WasbTaskHandler, "wasb_read")
    @mock.patch.object(WasbTaskHandler, "wasb_log_exists")
    def test_write_log(self, mock_log_exists, mock_wasb_read, mock_hook):
        mock_log_exists.return_value = True
        mock_wasb_read.return_value = ""
        self.wasb_task_handler.wasb_write('text', self.remote_log_location)
        mock_hook.return_value.load_string.assert_called_once_with(
            "text", self.container_name, self.remote_log_location)

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    @mock.patch.object(WasbTaskHandler, "wasb_read")
    @mock.patch.object(WasbTaskHandler, "wasb_log_exists")
    def test_write_on_existing_log(self, mock_log_exists, mock_wasb_read,
                                   mock_hook):
        mock_log_exists.return_value = True
        mock_wasb_read.return_value = "old log"
        self.wasb_task_handler.wasb_write('text', self.remote_log_location)
        mock_hook.return_value.load_string.assert_called_once_with(
            "old log\ntext", self.container_name, self.remote_log_location)

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_write_when_append_is_false(self, mock_hook):
        self.wasb_task_handler.wasb_write('text', self.remote_log_location,
                                          False)
        mock_hook.return_value.load_string.assert_called_once_with(
            "text", self.container_name, self.remote_log_location)

    def test_write_raises(self):
        handler = WasbTaskHandler(self.local_log_location,
                                  self.wasb_log_folder, self.container_name,
                                  self.filename_template, True)
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch(
                    "airflow.providers.microsoft.azure.hooks.wasb.WasbHook"
            ) as mock_hook:
                mock_hook.return_value.load_string.side_effect = AzureHttpError(
                    "failed to connect", 404)

                handler.wasb_write('text',
                                   self.remote_log_location,
                                   append=False)

            mock_error.assert_called_once_with('Could not write logs to %s',
                                               'remote/log/location/1.log',
                                               exc_info=True)
class TestWasbTaskHandler(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.wasb_log_folder = 'wasb://container/remote/log/location'
        self.remote_log_location = 'remote/log/location/1.log'
        self.local_log_location = 'local/log/location'
        self.container_name = "wasb-container"
        self.filename_template = '{try_number}.log'
        self.wasb_task_handler = WasbTaskHandler(
            base_log_folder=self.local_log_location,
            wasb_log_folder=self.wasb_log_folder,
            wasb_container=self.container_name,
            filename_template=self.filename_template,
            delete_local_copy=True,
        )

        date = datetime(2020, 8, 10)
        self.dag = DAG('dag_for_testing_file_task_handler', start_date=date)
        task = DummyOperator(task_id='task_for_testing_file_log_handler', dag=self.dag)
        self.ti = TaskInstance(task=task, execution_date=date)
        self.ti.try_number = 1
        self.ti.state = State.RUNNING
        self.addCleanup(self.dag.clear)

    @conf_vars({('logging', 'remote_log_conn_id'): 'wasb_default'})
    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.BlockBlobService")
    def test_hook(self, mock_service):
        assert isinstance(self.wasb_task_handler.hook, WasbHook)

    @conf_vars({('logging', 'remote_log_conn_id'): 'wasb_default'})
    def test_hook_raises(self):
        handler = WasbTaskHandler(
            self.local_log_location, self.wasb_log_folder, self.container_name, self.filename_template, True
        )
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") as mock_hook:
                mock_hook.side_effect = AzureHttpError("failed to connect", 404)
                # Initialize the hook
                handler.hook

            mock_error.assert_called_once_with(
                'Could not create an WasbHook with connection id "%s". '
                'Please make sure that airflow[azure] is installed and '
                'the Wasb connection exists.',
                "wasb_default",
            )

    def test_set_context_raw(self):
        self.ti.raw = True
        self.wasb_task_handler.set_context(self.ti)
        assert not self.wasb_task_handler.upload_on_close

    def test_set_context_not_raw(self):
        self.wasb_task_handler.set_context(self.ti)
        assert self.wasb_task_handler.upload_on_close

    # The `azure` provider uses legacy `azure-storage` library, where `snowflake` uses the
    # newer and more stable versions of those libraries. Most of `azure` operators and hooks work
    # fine together with `snowflake` because the deprecated library does not overlap with the
    # new libraries except the `blob` classes. So while `azure` works fine for most cases
    # blob is the only exception
    # Solution to that is being worked on in https://github.com/apache/airflow/pull/12188
    # Once this is merged, we can remove the xfail
    @pytest.mark.xfail
    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_wasb_log_exists(self, mock_hook):
        instance = mock_hook.return_value
        instance.check_for_blob.return_value = True
        self.wasb_task_handler.wasb_log_exists(self.remote_log_location)
        mock_hook.return_value.check_for_blob.assert_called_once_with(
            self.container_name, self.remote_log_location
        )

    # The `azure` provider uses legacy `azure-storage` library, where `snowflake` uses the
    # newer and more stable versions of those libraries. Most of `azure` operators and hooks work
    # fine together with `snowflake` because the deprecated library does not overlap with the
    # new libraries except the `blob` classes. So while `azure` works fine for most cases
    # blob is the only exception
    # Solution to that is being worked on in https://github.com/apache/airflow/pull/12188
    # Once this is merged, we can remove the xfail
    @pytest.mark.xfail
    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_wasb_read(self, mock_hook):
        mock_hook.return_value.read_file.return_value = 'Log line'
        assert self.wasb_task_handler.wasb_read(self.remote_log_location) == "Log line"
        assert self.wasb_task_handler.read(self.ti) == (
            [
                [
                    (
                        '',
                        '*** Reading remote log from wasb://container/remote/log/location/1.log.\n'
                        'Log line\n',
                    )
                ]
            ],
            [{'end_of_log': True}],
        )

    def test_wasb_read_raises(self):
        handler = WasbTaskHandler(
            self.local_log_location, self.wasb_log_folder, self.container_name, self.filename_template, True
        )
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") as mock_hook:
                mock_hook.return_value.read_file.side_effect = AzureHttpError("failed to connect", 404)

                handler.wasb_read(self.remote_log_location, return_error=True)

            mock_error.assert_called_once_with(
                'Could not read logs from remote/log/location/1.log', exc_info=True
            )

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    @mock.patch.object(WasbTaskHandler, "wasb_read")
    @mock.patch.object(WasbTaskHandler, "wasb_log_exists")
    def test_write_log(self, mock_log_exists, mock_wasb_read, mock_hook):
        mock_log_exists.return_value = True
        mock_wasb_read.return_value = ""
        self.wasb_task_handler.wasb_write('text', self.remote_log_location)
        mock_hook.return_value.load_string.assert_called_once_with(
            "text", self.container_name, self.remote_log_location
        )

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    @mock.patch.object(WasbTaskHandler, "wasb_read")
    @mock.patch.object(WasbTaskHandler, "wasb_log_exists")
    def test_write_on_existing_log(self, mock_log_exists, mock_wasb_read, mock_hook):
        mock_log_exists.return_value = True
        mock_wasb_read.return_value = "old log"
        self.wasb_task_handler.wasb_write('text', self.remote_log_location)
        mock_hook.return_value.load_string.assert_called_once_with(
            "old log\ntext", self.container_name, self.remote_log_location
        )

    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook")
    def test_write_when_append_is_false(self, mock_hook):
        self.wasb_task_handler.wasb_write('text', self.remote_log_location, False)
        mock_hook.return_value.load_string.assert_called_once_with(
            "text", self.container_name, self.remote_log_location
        )

    def test_write_raises(self):
        handler = WasbTaskHandler(
            self.local_log_location, self.wasb_log_folder, self.container_name, self.filename_template, True
        )
        with mock.patch.object(handler.log, 'error') as mock_error:
            with mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") as mock_hook:
                mock_hook.return_value.load_string.side_effect = AzureHttpError("failed to connect", 404)

                handler.wasb_write('text', self.remote_log_location, append=False)

            mock_error.assert_called_once_with(
                'Could not write logs to %s', 'remote/log/location/1.log', exc_info=True
            )