Subham-KRLX commented on code in PR #64051:
URL: https://github.com/apache/airflow/pull/64051#discussion_r3177742358
##########
providers/airbyte/tests/unit/airbyte/triggers/test_airbyte.py:
##########
@@ -36,6 +36,7 @@ class TestAirbyteSyncTrigger:
CONN_ID = "airbyte_default"
END_TIME = time.time() + 60 * 60 * 24 * 7
POLL_INTERVAL = 3.0
+ EXECUTION_DEADLINE = time.time() + 60 * 60 * 24 * 7
Review Comment:
time.time() at class level is evaluated once at import time the value will
be stale if tests run long after import could you move this into a fixture or
compute it inside the test?
##########
providers/airbyte/tests/unit/airbyte/operators/test_airbyte.py:
##########
@@ -68,8 +68,15 @@ def test_execute(self, mock_wait_for_job,
mock_submit_sync_connection, create_co
job_id=self.job_id, wait_seconds=self.wait_seconds,
timeout=self.timeout
)
- @pytest.mark.parametrize("status", ["success", "cancelled"])
- def test_execute_complete_non_error_states(self, status,
create_connection_without_db):
+ @pytest.mark.parametrize(
+ ("status", "should_raise", "expected_message"),
+ [
+ (JobStatusEnum.SUCCEEDED, False, "Job Succeeded"),
+ (JobStatusEnum.CANCELLED, True, "Job Cancelled"),
Review Comment:
CANCELLED shouldn't raise non-deferrable path treats it as a normal terminal
state this test is asserting the wrong behaviour.
##########
providers/airbyte/tests/unit/airbyte/triggers/test_airbyte.py:
##########
@@ -166,60 +164,106 @@ async def test_airbyte_job_for_terminal_status_error(
"message": mock_message,
"job_id": self.JOB_ID,
}
- task = asyncio.create_task(trigger.run().__anext__())
- await asyncio.sleep(0.5)
- assert TriggerEvent(expected_result) == task.result()
- asyncio.get_event_loop().stop()
+
+ events = [e async for e in trigger.run()]
+ assert len(events) == 1
+ assert TriggerEvent(expected_result) == events[0]
@pytest.mark.asyncio
-
@mock.patch("airflow.providers.airbyte.triggers.airbyte.AirbyteSyncTrigger.is_still_running")
@mock.patch("airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job_status")
- async def test_airbyte_job_exception(self, mock_get_job_status,
mocked_is_still_running):
+ async def test_airbyte_job_exception(self, mock_get_job_status):
"""Assert that run catch exception if Airbyte Sync job API throw
exception"""
- mocked_is_still_running.return_value = False
mock_get_job_status.side_effect = Exception("Test exception")
trigger = AirbyteSyncTrigger(
conn_id=self.CONN_ID,
poll_interval=self.POLL_INTERVAL,
end_time=self.END_TIME,
job_id=self.JOB_ID,
)
- task = [i async for i in trigger.run()]
- response = TriggerEvent(
+
+ events = [e async for e in trigger.run()]
+
+ expected_result = TriggerEvent(
{
"status": "error",
"message": "Test exception",
"job_id": self.JOB_ID,
}
)
- assert len(task) == 1
- assert response in task
+ assert len(events) == 1
+ assert expected_result in events
@pytest.mark.asyncio
-
@mock.patch("airflow.providers.airbyte.triggers.airbyte.AirbyteSyncTrigger.is_still_running")
@mock.patch("airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job_status")
- async def test_airbyte_job_timeout(self, mock_get_job_status,
mocked_is_still_running):
+ async def test_airbyte_job_timeout(self, mock_get_job_status):
"""Assert that run timeout after end_time elapsed"""
- mocked_is_still_running.return_value = True
mock_get_job_status.side_effect = Exception("Test exception")
Review Comment:
mock_get_job_status.side_effect is never hit here the deadline fires before
get_job_status is called can you remove it to keep the test setup clean?
##########
providers/airbyte/tests/unit/airbyte/triggers/test_airbyte.py:
##########
@@ -166,60 +164,106 @@ async def test_airbyte_job_for_terminal_status_error(
"message": mock_message,
"job_id": self.JOB_ID,
}
- task = asyncio.create_task(trigger.run().__anext__())
- await asyncio.sleep(0.5)
- assert TriggerEvent(expected_result) == task.result()
- asyncio.get_event_loop().stop()
+
+ events = [e async for e in trigger.run()]
+ assert len(events) == 1
+ assert TriggerEvent(expected_result) == events[0]
@pytest.mark.asyncio
-
@mock.patch("airflow.providers.airbyte.triggers.airbyte.AirbyteSyncTrigger.is_still_running")
@mock.patch("airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job_status")
- async def test_airbyte_job_exception(self, mock_get_job_status,
mocked_is_still_running):
+ async def test_airbyte_job_exception(self, mock_get_job_status):
"""Assert that run catch exception if Airbyte Sync job API throw
exception"""
- mocked_is_still_running.return_value = False
mock_get_job_status.side_effect = Exception("Test exception")
trigger = AirbyteSyncTrigger(
conn_id=self.CONN_ID,
poll_interval=self.POLL_INTERVAL,
end_time=self.END_TIME,
job_id=self.JOB_ID,
)
- task = [i async for i in trigger.run()]
- response = TriggerEvent(
+
+ events = [e async for e in trigger.run()]
+
+ expected_result = TriggerEvent(
{
"status": "error",
"message": "Test exception",
"job_id": self.JOB_ID,
}
)
- assert len(task) == 1
- assert response in task
+ assert len(events) == 1
+ assert expected_result in events
@pytest.mark.asyncio
-
@mock.patch("airflow.providers.airbyte.triggers.airbyte.AirbyteSyncTrigger.is_still_running")
@mock.patch("airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job_status")
- async def test_airbyte_job_timeout(self, mock_get_job_status,
mocked_is_still_running):
+ async def test_airbyte_job_timeout(self, mock_get_job_status):
"""Assert that run timeout after end_time elapsed"""
- mocked_is_still_running.return_value = True
mock_get_job_status.side_effect = Exception("Test exception")
+
end_time = time.time()
trigger = AirbyteSyncTrigger(
conn_id=self.CONN_ID,
poll_interval=self.POLL_INTERVAL,
end_time=end_time,
job_id=self.JOB_ID,
)
- generator = trigger.run()
- actual = await generator.asend(None)
- expected = TriggerEvent(
+
+ events = [e async for e in trigger.run()]
+
+ expected_result = TriggerEvent(
{
"status": "error",
"message": f"Job run {self.JOB_ID} has not reached a terminal
status "
f"after {end_time} seconds.",
"job_id": self.JOB_ID,
}
)
- assert expected == actual
+
+ assert len(events) == 1
+ assert expected_result in events
+
+ @pytest.mark.asyncio
+ async def test_airbyte_job_run_execution_timeout(self):
+ """Assert that run timeout after execution_deadline has elapsed"""
+
+ execution_deadline = time.time() - 1
+
+ trigger = AirbyteSyncTrigger(
+ conn_id=self.CONN_ID,
+ poll_interval=self.POLL_INTERVAL,
+ end_time=time.time(),
+ execution_deadline=execution_deadline,
+ job_id=self.JOB_ID,
+ )
Review Comment:
Both end_time and execution_deadline are already expired here so the test
doesn't verify which deadline fires first consider setting end_time to a future
value so the test specifically exercises the execution_deadline path.
##########
providers/airbyte/src/airflow/providers/airbyte/triggers/airbyte.py:
##########
@@ -70,6 +75,17 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
hook = AirbyteHook(airbyte_conn_id=self.conn_id)
try:
while await self.is_still_running(hook):
+ if self.execution_deadline is not None:
+ if self.execution_deadline < time.time():
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "message": f"Job run {self.job_id} has reached
execution timeout.",
+ "job_id": self.job_id,
+ }
+ )
+ return
+
if self.end_time < time.time():
yield TriggerEvent(
Review Comment:
Not resolved get_job_status is still called after the deadline check a job
that finishes after the deadline will get a timeout event instead of
success/cancelled.
##########
providers/airbyte/tests/unit/airbyte/operators/test_airbyte.py:
##########
@@ -126,9 +133,21 @@ def test_on_kill(self, mock_cancel_job,
mock_get_job_status, create_connection_w
connection_id=self.connection_id,
wait_seconds=self.wait_seconds,
timeout=self.timeout,
+ deferrable=True,
)
- op.job_id = self.job_id
- op.on_kill()
- mock_cancel_job.assert_called_once_with(self.job_id)
- mock_get_job_status.assert_called_once_with(self.job_id)
+ timeout_event = {
+ "status": "timeout",
+ "message": "Job run 1 has reached execution timeout.",
+ "job_id": self.job_id,
+ }
+
+ with pytest.raises(RuntimeError, match="has reached execution
timeout"):
Review Comment:
The PR description mentions that cancellation failures are logged and don't
mask the timeout error could you add a test for when cancel_job raises just to
make sure that behaviour is covered?
--
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]
