Repository: PrefectHQ/prefect
File: src/integrations/prefect-aws/tests/workers/test_ecs_worker.py (Python)

Code before:
                            "image": "{{ image }}",
                            "essential": False,
                        },
                        {
                            "name": "datadog-agent",
                            "essential": False,
                            "image": "datadog/agent",
                        },
                        {
                            "name": "log-router",
                            "essential": False,
                            "image": "public.ecr.aws/aws-observability/fluent-bit:latest",
                        },
                    ]
                }
            )
        )
    except ValueError as e:
        assert (
            "At least one container in the task definition must be marked as essential."
            in str(e)
        ), "Unexpected error message."
    else:
        assert False, (
            "ValueError was not raised when all containers were non-essential."
        )


@pytest.mark.usefixtures("ecs_mocks")
async def test_worker_cache_miss_for_registered_task_definitions_clears_from_cache(
    aws_credentials: AwsCredentials, flow_run: FlowRun
):
    configuration = await construct_configuration(
        aws_credentials=aws_credentials, command="echo test"
    )

    session = aws_credentials.get_boto3_session()
    ecs_client = session.client("ecs")

    async with ECSWorker(work_pool_name="test") as worker:
        result_1 = await worker.run(flow_run, configuration)

        # Fail to retrieve from cache on next run
        worker._retrieve_task_definition = MagicMock(
            side_effect=RuntimeError("failure retrieving from cache")
        )

        result_2 = await worker.run(flow_run, configuration)

    _, task_arn_1 = parse_identifier(result_1.identifier)
    task_1 = describe_task(ecs_client, task_arn_1)

Review comment (suggestion):
We prefer to use `pytest.raises` when asserting that exceptions are raised.

Code after:
                    "containerDefinitions": [
                        {
                            "name": ECS_DEFAULT_CONTAINER_NAME,
                            "image": "{{ image }}",
                            "essential": False,
                        },
                        {
                            "name": "datadog-agent",
                            "essential": False,
                            "image": "datadog/agent",
                        },
                        {
                            "name": "log-router",
                            "essential": False,
                            "image": "public.ecr.aws/aws-observability/fluent-bit:latest",
                        },
                    ]
                }
            )
        )


@pytest.mark.usefixtures("ecs_mocks")
async def test_worker_cache_miss_for_registered_task_definitions_clears_from_cache(
    aws_credentials: AwsCredentials, flow_run: FlowRun
):
    configuration = await construct_configuration(
        aws_credentials=aws_credentials, command="echo test"
    )

    session = aws_credentials.get_boto3_session()
    ecs_client = session.client("ecs")

    async with ECSWorker(work_pool_name="test") as worker:
        result_1 = await worker.run(flow_run, configuration)

        # Fail to retrieve from cache on next run
        worker._retrieve_task_definition = MagicMock(
            side_effect=RuntimeError("failure retrieving from cache")
        )

        result_2 = await worker.run(flow_run, configuration)

    _, task_arn_1 = parse_identifier(result_1.identifier)
    task_1 = describe_task(ecs_client, task_arn_1)
    _, task_arn_2 = parse_identifier(result_2.identifier)
    task_2 = describe_task(ecs_client, task_arn_2)

    assert task_1["taskDefinitionArn"] != task_2["taskDefinitionArn"]
    assert task_1["taskDefinitionArn"] not in _TASK_DEFINITION_CACHE.values(), (
        _TASK_DEFINITION_CACHE