o-nikolas commented on code in PR #64745:
URL: https://github.com/apache/airflow/pull/64745#discussion_r3455583616
##########
providers/amazon/tests/unit/amazon/aws/operators/test_batch.py:
##########
@@ -545,46 +551,66 @@ def test_monitor_job_with_logs(
@patch.object(BatchOperator, "log", new_callable=MagicMock)
@patch("airflow.providers.amazon.aws.operators.batch.validate_execute_complete_event")
- @patch.object(BatchOperator, "monitor_job")
- def test_execute_complete_success_with_logs(self, mock_monitor_job,
mock_validate, mock_log):
+ @patch.object(BatchClientHook, "check_job_success")
+ @patch.object(BatchClientHook, "get_job_all_awslogs_info")
+
@mock.patch("airflow.providers.amazon.aws.hooks.batch_client.AwsBaseHook.get_client_type")
+ def test_execute_complete_success_with_logs(
+ self, mock_client, mock_get_job_all_awslogs_info,
mock_check_job_success, mock_validate, mock_log
+ ):
# Setup
mock_validate.return_value = {"status": "success", "job_id": "12345"}
+ mock_get_job_all_awslogs_info.return_value = [
+ {"awslogs_group": "/aws/batch/job", "awslogs_stream_name":
"stream1"}
+ ]
batch = BatchOperator(
task_id="test_task",
job_name=JOB_NAME,
job_queue="dummy_queue",
job_definition="dummy_definition",
deferrable=True,
awslogs_enabled=True,
+ do_xcom_push=True, # Enable xcom push so _persist_cloudwatch_link
runs
)
+ # Add task instance to context
+ context = {"ti": mock.MagicMock()}
- result = batch.execute_complete(context={}, event={"dummy": "event"})
+ result = batch.execute_complete(context=context, event={"dummy":
"event"})
# Assertion
assert result == "12345"
- mock_monitor_job.assert_called_once_with({})
+ mock_get_job_all_awslogs_info.assert_called_once_with("12345")
+ mock_check_job_success.assert_called_once_with("12345")
mock_log.info.assert_called_with("Job completed successfully for
job_id: %s", "12345")
@patch.object(BatchOperator, "log", new_callable=MagicMock)
@patch("airflow.providers.amazon.aws.operators.batch.validate_execute_complete_event")
- @patch.object(BatchOperator, "monitor_job")
- def test_execute_complete_success_without_logs(self, mock_monitor_job,
mock_validate, mock_log):
+ @patch.object(BatchClientHook, "check_job_success")
+ @patch.object(BatchClientHook, "get_job_all_awslogs_info")
+
@mock.patch("airflow.providers.amazon.aws.hooks.batch_client.AwsBaseHook.get_client_type")
+ def test_execute_complete_success_without_logs(
Review Comment:
We're missing testing of the failure case no? To ensure the log links are
still persisted despite failures?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py:
##########
@@ -252,14 +258,17 @@ def execute(self, context: Context) -> str | None:
def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> str:
validated_event = validate_execute_complete_event(event)
+ # Set job_id first so CloudWatch link can be persisted even on failure
+ self.job_id = validated_event["job_id"]
+
+ # Persist CloudWatch logs for both success and failure
+ self._persist_cloudwatch_link(context)
+
if validated_event["status"] != "success":
raise AirflowException(f"Error while running job:
{validated_event}")
- self.job_id = validated_event["job_id"]
-
- # Fetch logs if awslogs_enabled
- if self.awslogs_enabled:
- self.monitor_job(context) # fetch logs, no need to return
Review Comment:
It looks like this just got dropped in the new changes? How are logs pulled
down now?
--
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]
