From a5743ce5f369b1e3ef4602ba7a10d26f3e738650 Mon Sep 17 00:00:00 2001
From: jeanluciano <jeanluciano@mailbox.org>
Date: Wed, 10 Apr 2024 10:49:40 -0500
Subject: [PATCH] mock run task call

---
 prefect_aws/workers/ecs_worker.py |  1 +
 tests/workers/test_ecs_worker.py  | 37 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/prefect_aws/workers/ecs_worker.py b/prefect_aws/workers/ecs_worker.py
index a76dc49b..da783c8a 100644
--- a/prefect_aws/workers/ecs_worker.py
+++ b/prefect_aws/workers/ecs_worker.py
@@ -1567,6 +1567,7 @@ def _prepare_task_run_request(
             CREATE_TASK_RUN_MIN_DELAY_JITTER_SECONDS,
             CREATE_TASK_RUN_MAX_DELAY_JITTER_SECONDS,
         ),
+        reraise=True,
     )
     def _create_task_run(self, ecs_client: _ECSClient, task_run_request: dict) -> str:
         """
diff --git a/tests/workers/test_ecs_worker.py b/tests/workers/test_ecs_worker.py
index 774e4692..6fbf3fb0 100644
--- a/tests/workers/test_ecs_worker.py
+++ b/tests/workers/test_ecs_worker.py
@@ -3,9 +3,11 @@
 from functools import partial
 from typing import Any, Awaitable, Callable, Dict, List, Optional
 from unittest.mock import ANY, MagicMock
+from unittest.mock import patch as mock_patch
 from uuid import uuid4
 
 import anyio
+import botocore
 import pytest
 import yaml
 from moto import mock_ec2, mock_ecs, mock_logs
@@ -1322,6 +1324,41 @@ async def write_fake_log(task_arn):
     # assert "test-message-{i}" in err
 
 
+orig = botocore.client.BaseClient._make_api_call
+
+
+def mock_make_api_call(self, operation_name, kwarg):
+    if operation_name == "RunTask":
+        return {
+            "failures": [
+                {"arn": "string", "reason": "string", "detail": "string"},
+            ]
+        }
+    return orig(self, operation_name, kwarg)
+
+
+@pytest.mark.usefixtures("ecs_mocks")
+async def test_run_task_error_handling(
+    aws_credentials: AwsCredentials,
+    flow_run: FlowRun,
+    capsys,
+):
+    configuration = await construct_configuration(
+        aws_credentials=aws_credentials,
+        task_role_arn="test",
+    )
+
+    with mock_patch(
+        "botocore.client.BaseClient._make_api_call", new=mock_make_api_call
+    ):
+        async with ECSWorker(work_pool_name="test") as worker:
+            with pytest.raises(RuntimeError) as exc:
+                await run_then_stop_task(worker, configuration, flow_run)
+
+    print(exc.value)
+    assert "Failed to run ECS task: string" in capsys.readouterr().out
+
+
 @pytest.mark.usefixtures("ecs_mocks")
 async def test_cloudwatch_log_options(
     aws_credentials: AwsCredentials, flow_run: FlowRun