subhramit commented on code in PR #73522:
URL: https://github.com/apache/airflow/pull/73522#discussion_r4088215388
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -79,3 +79,111 @@ async def test_run_propagates_cancellation(self,
mock_hook_class):
with pytest.raises(asyncio.CancelledError):
await anext(trigger.run())
+
+
+class TestInfluxDB3SensorTrigger:
+ def test_serialization(self):
+ trigger = InfluxDB3SensorTrigger(
+ sql=SQL,
+ influxdb3_conn_id=CONN_ID,
+ poll_interval=30,
+ fail_on_empty=True,
+ )
+
+ assert trigger.serialize() == (
+
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3SensorTrigger",
+ {
+ "sql": SQL,
+ "influxdb3_conn_id": CONN_ID,
+ "poll_interval": 30,
+ "fail_on_empty": True,
+ },
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_succeeds_without_sleep_when_condition_is_met(self,
mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": [1]}))
+
+ events = [event async for event in
InfluxDB3SensorTrigger(sql=SQL).run()]
+
+ mock_hook_class.assert_called_once_with(conn_id="influxdb3_default")
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_not_awaited()
+ assert events == [TriggerEvent({"status": "success"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_polls_until_condition_is_met(self, mock_hook_class,
mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(
+ side_effect=[
+ pd.DataFrame({"literal": []}),
+ pd.DataFrame({"count": [0]}),
+ pd.DataFrame({"literal": [1]}),
+ ]
+ )
+
+ events = [event async for event in InfluxDB3SensorTrigger(sql=SQL,
poll_interval=30).run()]
+
+ assert mock_hook.query_async.await_count == 3
+ assert mock_sleep.await_args_list == [mock.call(30), mock.call(30)]
+ assert events == [TriggerEvent({"status": "success"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_fail_on_empty(self, mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": []}))
+
+ events = [event async for event in InfluxDB3SensorTrigger(sql=SQL,
fail_on_empty=True).run()]
+
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_not_awaited()
+ assert events == [
+ TriggerEvent(
+ {"status": "error", "message": "No rows returned, raising as
per fail_on_empty flag"}
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_failure(self, mock_hook_class):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(side_effect=ValueError("boom"))
+
+ events = [event async for event in
InfluxDB3SensorTrigger(sql=SQL).run()]
+
+ mock_hook_class.assert_called_once_with(conn_id="influxdb3_default")
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ assert events == [TriggerEvent({"status": "error", "message": "boom"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_propagates_cancellation_during_query(self,
mock_hook_class):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(side_effect=asyncio.CancelledError())
+
+ with pytest.raises(asyncio.CancelledError):
+ await anext(InfluxDB3SensorTrigger(sql=SQL).run())
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
+ new_callable=mock.AsyncMock,
+ side_effect=asyncio.CancelledError(),
+ )
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_propagates_cancellation_during_sleep(self,
mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": []}))
+
+ with pytest.raises(asyncio.CancelledError):
+ await anext(InfluxDB3SensorTrigger(sql=SQL,
poll_interval=30).run())
+
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_awaited_once_with(30)
Review Comment:
Added in
https://github.com/apache/airflow/pull/73522/commits/2ce4be1625c6f74430631a7a4de097b5937e707f,
thanks. That makes sense. Neither me nor my agent could think of that!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
