Fix WasbPrefixSensor arg inconsistency between sync and async mode

airflow/providers/microsoft/azure/sensors/wasb.py

@@ -149,6 +149,8 @@ class WasbPrefixSensor(BaseSensorOperator):
     :param wasb_conn_id: Reference to the wasb connection.
     :param check_options: Optional keyword arguments that
         `WasbHook.check_for_prefix()` takes.
+    :param public_read: whether an anonymous public read access should be used. Default is False
+    :param deferrable: Run operator in the deferrable mode.
     """
 
     template_fields: Sequence[str] = ("container_name", "prefix")
@@ -160,21 +162,23 @@ def __init__(
         prefix: str,
         wasb_conn_id: str = "wasb_default",
         check_options: dict | None = None,
+        public_read: bool = False,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         if check_options is None:
             check_options = {}
-        self.wasb_conn_id = wasb_conn_id
         self.container_name = container_name
         self.prefix = prefix
+        self.wasb_conn_id = wasb_conn_id
         self.check_options = check_options
+        self.public_read = public_read
         self.deferrable = deferrable
 
     def poke(self, context: Context) -> bool:
         self.log.info("Poking for prefix: %s in wasb://%s", self.prefix, self.container_name)
-        hook = WasbHook(wasb_conn_id=self.wasb_conn_id)
+        hook = WasbHook(wasb_conn_id=self.wasb_conn_id, public_read=self.public_read)
         return hook.check_for_prefix(self.container_name, self.prefix, **self.check_options)
 
     def execute(self, context: Context) -> None:
@@ -193,6 +197,8 @@ def execute(self, context: Context) -> None:
                         container_name=self.container_name,
                         prefix=self.prefix,
                         wasb_conn_id=self.wasb_conn_id,
+                        check_options=self.check_options,
+                        public_read=self.public_read,
                         poke_interval=self.poke_interval,
                     ),
                     method_name="execute_complete",

airflow/providers/microsoft/azure/triggers/wasb.py

@@ -102,26 +102,28 @@ class WasbPrefixSensorTrigger(BaseTrigger):
             ``copy``, ``deleted``
     :param delimiter: filters objects based on the delimiter (for e.g '.csv')
     :param wasb_conn_id: the connection identifier for connecting to Azure WASB
-    :param poke_interval:  polling period in seconds to check for the status
+    :param check_options: Optional keyword arguments that
+        `WasbAsyncHook.check_for_prefix_async()` takes.
     :param public_read: whether an anonymous public read access should be used. Default is False
+    :param poke_interval:  polling period in seconds to check for the status
     """
 
     def __init__(
         self,
         container_name: str,
         prefix: str,
-        include: list[str] | None = None,
-        delimiter: str = "/",
         wasb_conn_id: str = "wasb_default",
+        check_options: dict | None = None,
         public_read: bool = False,
         poke_interval: float = 5.0,
     ):
+        if not check_options:
+            check_options = {}
         super().__init__()
         self.container_name = container_name
         self.prefix = prefix
-        self.include = include
-        self.delimiter = delimiter
         self.wasb_conn_id = wasb_conn_id
+        self.check_options = check_options
         self.poke_interval = poke_interval
         self.public_read = public_read
 
@@ -132,10 +134,9 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             {
                 "container_name": self.container_name,
                 "prefix": self.prefix,
-                "include": self.include,
-                "delimiter": self.delimiter,
                 "wasb_conn_id": self.wasb_conn_id,
                 "poke_interval": self.poke_interval,
+                "check_options": self.check_options,
                 "public_read": self.public_read,
             },
         )
@@ -148,10 +149,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
             async with await hook.get_async_conn():
                 while not prefix_exists:
                     prefix_exists = await hook.check_for_prefix_async(
-                        container_name=self.container_name,
-                        prefix=self.prefix,
-                        include=self.include,
-                        delimiter=self.delimiter,
+                        container_name=self.container_name, prefix=self.prefix, **self.check_options
                     )
                     if prefix_exists:
                         message = f"Prefix {self.prefix} found in container {self.container_name}."

tests/always/test_project_structure.py

@@ -158,7 +158,6 @@ def test_providers_modules_should_have_tests(self):
             "tests/providers/microsoft/azure/fs/test_adls.py",
             "tests/providers/microsoft/azure/operators/test_adls.py",
             "tests/providers/microsoft/azure/transfers/test_azure_blob_to_gcs.py",
-            "tests/providers/microsoft/azure/triggers/test_wasb.py",
             "tests/providers/mongo/sensors/test_mongo.py",
             "tests/providers/openlineage/extractors/test_manager.py",
             "tests/providers/openlineage/plugins/test_adapter.py",

tests/providers/microsoft/azure/triggers/test_wasb.py

@@ -0,0 +1,220 @@
+# 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.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from unittest import mock
+
+import pytest
+
+from airflow.providers.microsoft.azure.triggers.wasb import (
+    WasbBlobSensorTrigger,
+    WasbPrefixSensorTrigger,
+)
+from airflow.triggers.base import TriggerEvent
+
+TEST_DATA_STORAGE_BLOB_NAME = "test_blob_providers_team.txt"
+TEST_DATA_STORAGE_CONTAINER_NAME = "test-container-providers-team"
+TEST_DATA_STORAGE_BLOB_PREFIX = TEST_DATA_STORAGE_BLOB_NAME[:10]
+TEST_WASB_CONN_ID = "wasb_default"
+POKE_INTERVAL = 5.0
+
+
+class TestWasbBlobSensorTrigger:
+    TRIGGER = WasbBlobSensorTrigger(
+        container_name=TEST_DATA_STORAGE_CONTAINER_NAME,
+        blob_name=TEST_DATA_STORAGE_BLOB_NAME,
+        wasb_conn_id=TEST_WASB_CONN_ID,
+        poke_interval=POKE_INTERVAL,
+    )
+
+    def test_serialization(self):
+        """
+        Asserts that the WasbBlobSensorTrigger correctly serializes its arguments
+        and classpath.
+        """
+
+        classpath, kwargs = self.TRIGGER.serialize()
+        assert classpath == "airflow.providers.microsoft.azure.triggers.wasb.WasbBlobSensorTrigger"
+        assert kwargs == {
+            "container_name": TEST_DATA_STORAGE_CONTAINER_NAME,
+            "blob_name": TEST_DATA_STORAGE_BLOB_NAME,
+            "wasb_conn_id": TEST_WASB_CONN_ID,
+            "poke_interval": POKE_INTERVAL,
+            "public_read": False,
+        }
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "blob_exists",
+        [True, False],
+    )
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_blob_async")
+    async def test_running(self, mock_check_for_blob, blob_exists):
+        """
+        Test if the task is run in trigger successfully.
+        """
+        mock_check_for_blob.return_value = blob_exists
+
+        task = asyncio.create_task(self.TRIGGER.run().__anext__())
+
+        # TriggerEvent was not returned
+        assert task.done() is False
+        asyncio.get_event_loop().stop()
+
+    @pytest.mark.db_test
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_blob_async")
+    async def test_success(self, mock_check_for_blob):
+        """Tests the success state for that the WasbBlobSensorTrigger."""
+        mock_check_for_blob.return_value = True
+
+        task = asyncio.create_task(self.TRIGGER.run().__anext__())
+        await asyncio.sleep(0.5)
+
+        # TriggerEvent was returned
+        assert task.done() is True
+        asyncio.get_event_loop().stop()
+
+        message = f"Blob {TEST_DATA_STORAGE_BLOB_NAME} found in container {TEST_DATA_STORAGE_CONTAINER_NAME}."
+        assert task.result() == TriggerEvent({"status": "success", "message": message})
+
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_blob_async")
+    async def test_waiting_for_blob(self, mock_check_for_blob, caplog):
+        """Tests the WasbBlobSensorTrigger sleeps waiting for the blob to arrive."""
+        mock_check_for_blob.side_effect = [False, True]
+        caplog.set_level(logging.INFO)
+
+        with mock.patch.object(self.TRIGGER.log, "info"):
+            task = asyncio.create_task(self.TRIGGER.run().__anext__())
+
+        await asyncio.sleep(POKE_INTERVAL + 0.5)
+
+        if not task.done():
+            message = (
+                f"Blob {TEST_DATA_STORAGE_BLOB_NAME} not available yet in container {TEST_DATA_STORAGE_CONTAINER_NAME}."
+                f" Sleeping for {POKE_INTERVAL} seconds"
+            )
+            assert message in caplog.text
+        asyncio.get_event_loop().stop()
+
+    @pytest.mark.db_test
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_blob_async")
+    async def test_trigger_exception(self, mock_check_for_blob):
+        """Tests the WasbBlobSensorTrigger yields an error event if there is an exception."""
+        mock_check_for_blob.side_effect = Exception("Test exception")
+
+        task = [i async for i in self.TRIGGER.run()]
+        assert len(task) == 1
+        assert TriggerEvent({"status": "error", "message": "Test exception"}) in task
+
+
+class TestWasbPrefixSensorTrigger:
+    TRIGGER = WasbPrefixSensorTrigger(
+        container_name=TEST_DATA_STORAGE_CONTAINER_NAME,
+        prefix=TEST_DATA_STORAGE_BLOB_PREFIX,
+        wasb_conn_id=TEST_WASB_CONN_ID,
+        poke_interval=POKE_INTERVAL,
+        check_options={"delimiter": "/", "include": None},
+    )
+
+    def test_serialization(self):
+        """
+        Asserts that the WasbPrefixSensorTrigger correctly serializes its arguments and classpath."""
+
+        classpath, kwargs = self.TRIGGER.serialize()
+        assert classpath == "airflow.providers.microsoft.azure.triggers.wasb.WasbPrefixSensorTrigger"
+        assert kwargs == {
+            "container_name": TEST_DATA_STORAGE_CONTAINER_NAME,
+            "prefix": TEST_DATA_STORAGE_BLOB_PREFIX,
+            "wasb_conn_id": TEST_WASB_CONN_ID,
+            "public_read": False,
+            "check_options": {
+                "delimiter": "/",
+                "include": None,
+            },
+            "poke_interval": POKE_INTERVAL,
+        }
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "prefix_exists",
+        [True, False],
+    )
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_prefix_async")
+    async def test_running(self, mock_check_for_prefix, prefix_exists):
+        """Test if the task is run in trigger successfully."""
+        mock_check_for_prefix.return_value = prefix_exists
+
+        task = asyncio.create_task(self.TRIGGER.run().__anext__())
+
+        # TriggerEvent was not returned
+        assert task.done() is False
+        asyncio.get_event_loop().stop()
+
+    @pytest.mark.db_test
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_prefix_async")
+    async def test_success(self, mock_check_for_prefix):
+        """Tests the success state for that the WasbPrefixSensorTrigger."""
+        mock_check_for_prefix.return_value = True
+
+        task = asyncio.create_task(self.TRIGGER.run().__anext__())
+        await asyncio.sleep(0.5)
+
+        # TriggerEvent was returned
+        assert task.done() is True
+        asyncio.get_event_loop().stop()
+
+        message = (
+            f"Prefix {TEST_DATA_STORAGE_BLOB_PREFIX} found in container {TEST_DATA_STORAGE_CONTAINER_NAME}."
+        )
+        assert task.result() == TriggerEvent({"status": "success", "message": message})
+
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_prefix_async")
+    async def test_waiting_for_blob(self, mock_check_for_prefix):
+        """Tests the WasbPrefixSensorTrigger sleeps waiting for the blob to arrive."""
+        mock_check_for_prefix.side_effect = [False, True]
+
+        with mock.patch.object(self.TRIGGER.log, "info") as mock_log_info:
+            task = asyncio.create_task(self.TRIGGER.run().__anext__())
+
+        await asyncio.sleep(POKE_INTERVAL + 0.5)
+
+        if not task.done():
+            message = (
+                f"Prefix {TEST_DATA_STORAGE_BLOB_PREFIX} not available yet in container "
+                f"{TEST_DATA_STORAGE_CONTAINER_NAME}. Sleeping for {POKE_INTERVAL} seconds"
+            )
+            mock_log_info.assert_called_once_with(message)
+        asyncio.get_event_loop().stop()
+
+    @pytest.mark.db_test
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbAsyncHook.check_for_prefix_async")
+    async def test_trigger_exception(self, mock_check_for_prefix):
+        """Tests the WasbPrefixSensorTrigger yields an error event if there is an exception."""
+        mock_check_for_prefix.side_effect = Exception("Test exception")
+
+        task = [i async for i in self.TRIGGER.run()]
+        assert len(task) == 1
+        assert TriggerEvent({"status": "error", "message": "Test exception"}) in task