def test_parse_task_run_kwargs(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() kwarg_dict = { "cluster": "test", "count": "test", "startedBy": "test", "group": "test", "placementConstraints": "test", "placementStrategy": "test", "platformVersion": "test", "networkConfiguration": "test", "tags": "test", "enableECSManagedTags": "test", "propagateTags": "test", } task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict) assert task_run_kwargs == kwarg_dict assert task_definition_kwargs == { "placementConstraints": "test", "tags": "test" }
def test_fargate_agent_start_max_polls_zero(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) on_shutdown = MagicMock() monkeypatch.setattr("prefect.agent.fargate.agent.FargateAgent.on_shutdown", on_shutdown) agent_process = MagicMock() monkeypatch.setattr("prefect.agent.agent.Agent.agent_process", agent_process) agent_connect = MagicMock(return_value="id") monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect", agent_connect) heartbeat = MagicMock() monkeypatch.setattr("prefect.agent.fargate.agent.FargateAgent.heartbeat", heartbeat) agent = FargateAgent(max_polls=0) agent.start() assert on_shutdown.call_count == 1 assert agent_process.call_count == 0 assert heartbeat.call_count == 0
def test_deploy_flows_register_task_definition(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.run_task.return_value = {} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent() agent.deploy_flows(flow_runs=[ GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", }) ]) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert (boto3_client.register_task_definition.call_args[1]["family"] == "prefect-task-id")
def test_deploy_flows_require_docker_storage(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {"tags": []} boto3_client.run_task.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) with pytest.raises(Exception) as excinfo: agent = FargateAgent() agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Local().serialize(), "id": "id", "version": 2, "name": "name", }), "id": "id", "name": "name", })) assert boto3_client.describe_task_definition.not_called assert boto3_client.run_task.not_called
def test_deploy_flow_raises(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {} boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent() agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", "name": "name", })) assert boto3_client.describe_task_definition.called assert boto3_client.run_task.called
def test_deploy_flows_no_security_group(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {} boto3_client.run_task.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent() agent.deploy_flows(flow_runs=[ GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", }) ]) assert boto3_client.describe_task_definition.called assert boto3_client.run_task.called assert boto3_client.run_task.call_args[1]["cluster"] == "default" assert boto3_client.run_task.call_args[1]["networkConfiguration"] == { "awsvpcConfiguration": { "subnets": [], "assignPublicIp": "ENABLED" } }
def test_deploy_flow_register_task_definition_no_repo_credentials( monkeypatch, runner_token ): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) with set_temporary_config({"logging.log_to_cloud": True}): agent = FargateAgent() agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": Docker( registry_url="test", image_name="name", image_tag="tag" ).serialize(), "id": "id", } ), "id": "id", } ) ) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert boto3_client.register_task_definition.call_args[1][ "containerDefinitions" ] == [ { "name": "flow", "image": "test/name:tag", "command": ["/bin/sh", "-c", "prefect execute cloud-flow"], "environment": [ {"name": "PREFECT__CLOUD__API", "value": "https://api.prefect.io"}, {"name": "PREFECT__CLOUD__AGENT__LABELS", "value": "[]"}, {"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"}, {"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"}, {"name": "PREFECT__LOGGING__LEVEL", "value": "DEBUG"}, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, ], "essential": True, } ]
def test_parse_task_kwargs_invalid_value_removed(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() kwarg_dict = {"test": "not_real"} task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict) assert task_definition_kwargs == {} assert task_run_kwargs == {}
def test_deploy_flows_enable_task_revisions_old_version_exists( monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = { "tags": [ { "key": "PrefectFlowId", "value": "current_id" }, { "key": "PrefectFlowVersion", "value": "5" }, ] } boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.get_resources.return_value = { "ResourceTagMappingList": [{ "ResourceARN": "arn:aws:ecs:us-east-1:12345:task-definition/flow:22" }] } monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent(enable_task_revisions=True) agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", "version": 3, "name": "name", }), "id": "id", "name": "name", })) assert boto3_client.describe_task_definition.called assert boto3_client.get_resources.called assert boto3_client.register_task_definition.not_called assert boto3_client.run_task.called assert (agent.task_definition_name == "arn:aws:ecs:us-east-1:12345:task-definition/flow:22")
def test_override_kwargs(monkeypatch, runner_token): boto3_resource = MagicMock() streaming_body = MagicMock() streaming_body.read.return_value.decode.return_value = ( '{"cpu": "256", "networkConfiguration": "test"}') boto3_resource.return_value.Object.return_value.get.return_value = { "Body": streaming_body } monkeypatch.setattr("boto3.resource", boto3_resource) agent = FargateAgent( use_external_kwargs=True, external_kwargs_s3_bucket="test-bucket", external_kwargs_s3_key="prefect-artifacts/kwargs", aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", region_name="region", labels=["aws", "staging"], ) definition_kwargs = {} run_kwargs = {} agent._override_kwargs( GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", "version": 2, "name": "name", }), "id": "id", "name": "name", }), definition_kwargs, run_kwargs, ) assert boto3_resource.called assert streaming_body.read().decode.called assert definition_kwargs == {"cpu": "256"} assert run_kwargs == {"networkConfiguration": "test"}
def test_ecs_agent_config_env_vars(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") monkeypatch.setenv("AWS_SESSION_TOKEN", "token") monkeypatch.setenv("TASK_ROLE_ARN", "task_role_arn") monkeypatch.setenv("EXECUTION_ROLE_ARN", "execution_role_arn") monkeypatch.setenv("REGION_NAME", "region") monkeypatch.setenv("CLUSTER", "cluster") monkeypatch.setenv("REPOSITORY_CREDENTIALS", "repo") monkeypatch.setenv("ASSIGN_PUBLIC_IP", "DISABLED") monkeypatch.setenv("TASK_CPU", "1") monkeypatch.setenv("TASK_MEMORY", "2") agent = FargateAgent(subnets=["subnet"]) assert agent assert agent.execution_role_arn == "execution_role_arn" assert agent.task_role_arn == "task_role_arn" assert agent.cluster == "cluster" assert agent.repository_credentials == "repo" assert agent.assign_public_ip == "DISABLED" assert agent.task_cpu == "1" assert agent.task_memory == "2" boto3_client.assert_called_with( "ecs", aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", region_name="region", )
def test_fargate_agent_config_options(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) # Client args monkeypatch.setenv("AWS_ACCESS_KEY_ID", "") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "") monkeypatch.setenv("AWS_SESSION_TOKEN", "") monkeypatch.setenv("REGION_NAME", "") monkeypatch.delenv("AWS_ACCESS_KEY_ID") monkeypatch.delenv("AWS_SECRET_ACCESS_KEY") monkeypatch.delenv("AWS_SESSION_TOKEN") monkeypatch.delenv("REGION_NAME") with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = FargateAgent(name="test", labels=["test"]) assert agent assert agent.labels == ["test"] assert agent.name == "test" assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert agent.boto3_client boto3_client.assert_called_with( "ecs", aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, region_name=None, )
def test_ecs_agent_init(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() assert agent assert agent.boto3_client
def test_ecs_agent_config_options_init(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent( aws_access_key_id="id", aws_secret_access_key="secret", region_name="region", cluster="cluster", subnets=["subnet"], security_groups=["security_group"], repository_credentials="repo", assign_public_ip="DISABLED", task_cpu="1", task_memory="2", ) assert agent assert agent.cluster == "cluster" assert agent.subnets == ["subnet"] assert agent.security_groups == ["security_group"] assert agent.repository_credentials == "repo" assert agent.assign_public_ip == "DISABLED" assert agent.task_cpu == "1" assert agent.task_memory == "2" boto3_client.assert_called_with( "ecs", aws_access_key_id="id", aws_secret_access_key="secret", region_name="region", )
def test_parse_task_definition_kwargs_errors(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() kwarg_dict = { "placementConstraints": "taskRoleArn='arn:aws:iam::543216789012:role/Dev", } task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict) assert task_definition_kwargs == kwarg_dict assert task_run_kwargs == { "placementConstraints": "taskRoleArn='arn:aws:iam::543216789012:role/Dev" }
def test_deploy_flows_enable_task_revisions_tags_passed_in(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = { "tags": [{"key": "PrefectFlowId", "value": "id"}] } boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent( enable_task_revisions=True, tags=[ {"key": "PrefectFlowId", "value": "id"}, {"key": "PrefectFlowVersion", "value": "2"}, ], ) agent.deploy_flow( GraphQLResult( { "flow": GraphQLResult( { "storage": Docker( registry_url="test", image_name="name", image_tag="tag" ).serialize(), "id": "id", "version": 2, "name": "name", } ), "id": "id", "name": "name", } ) ) assert agent.task_definition_kwargs == { "tags": [ {"key": "PrefectFlowId", "value": "id"}, {"key": "PrefectFlowVersion", "value": "2"}, ] } assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.not_called assert boto3_client.run_task.called assert boto3_client.run_task.called_with(taskDefinition="name")
def test_deploy_flow_register_task_definition_uses_user_env_vars( monkeypatch, runner_token ): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent(env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar")) agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": Docker( registry_url="test", image_name="name", image_tag="tag" ).serialize(), "id": "id", } ), "id": "id", } ) ) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert ( boto3_client.register_task_definition.call_args[1]["family"] == "prefect-task-id" ) container_defs = boto3_client.register_task_definition.call_args[1][ "containerDefinitions" ] user_vars = [ dict(name="AUTH_THING", value="foo"), dict(name="PKG_SETTING", value="bar"), ] assert container_defs[0]["environment"][-1] in user_vars assert container_defs[0]["environment"][-2] in user_vars
def test_fargate_agent_config_options_default(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() assert agent assert agent.labels == [] assert agent.name == "agent" assert agent.task_definition_kwargs == {} assert agent.task_run_kwargs == {} assert agent.boto3_client
def test_default_subnets(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_subnets.return_value = { "Subnets": [ {"MapPublicIpOnLaunch": False, "SubnetId": "id"}, {"MapPublicIpOnLaunch": True, "SubnetId": "id2"}, ] } monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent() assert agent.subnets == ["id"]
def test_k8s_agent_config_options(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = FargateAgent(name="test", labels=["test"]) assert agent assert agent.labels == ["test"] assert agent.name == "test" assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert agent.boto3_client
def test_deploy_flow_raises(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {} boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent() with pytest.raises(ValueError): agent.deploy_flow(flow_run=GraphQLResult( { "flow": GraphQLResult({ "storage": Local().serialize(), "id": "id" }), "id": "id", })) assert not boto3_client.describe_task_definition.called assert not boto3_client.run_task.called
def test_ecs_agent_config_options_default(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() assert agent assert agent.cluster == "default" assert not agent.subnets assert not agent.security_groups assert not agent.repository_credentials assert agent.assign_public_ip == "ENABLED" assert agent.task_cpu == "256" assert agent.task_memory == "512" assert agent.boto3_client
def test_parse_task_definition_kwargs(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) agent = FargateAgent() kwarg_dict = { "taskRoleArn": "test", "executionRoleArn": "test", "volumes": "test", "placementConstraints": "test", "cpu": "test", "memory": "test", "tags": "test", "pidMode": "test", "ipcMode": "test", "proxyConfiguration": "test", "inferenceAccelerators": "test", } task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict) assert task_definition_kwargs == kwarg_dict assert task_run_kwargs == {"placementConstraints": "test", "tags": "test"}
def test_fargate_agent_config_env_vars_lists_dicts(monkeypatch, runner_token): boto3_client = MagicMock() monkeypatch.setattr("boto3.client", boto3_client) def_kwarg_dict = { "placementConstraints": ["test"], "proxyConfiguration": { "test": "test" }, } run_kwarg_dict = { "placementConstraints": ["test"], "networkConfiguration": { "test": "test" }, } # Client args monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") monkeypatch.setenv("AWS_SESSION_TOKEN", "token") monkeypatch.setenv("REGION_NAME", "region") # Def / run args monkeypatch.setenv("placementConstraints", "['test']") monkeypatch.setenv("proxyConfiguration", "{'test': 'test'}") monkeypatch.setenv("networkConfiguration", "{'test': 'test'}") agent = FargateAgent(subnets=["subnet"]) assert agent assert agent.task_definition_kwargs == def_kwarg_dict assert agent.task_run_kwargs == run_kwarg_dict boto3_client.assert_called_with( "ecs", aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", region_name="region", )
def get_fargate_agent_definition(): """ Get the Prefect Agent definition for an environment that run workflows on AWS ECS Fargate Returns: [FargateAgent] -- Fargate Agent object from prefect.agent.fargate """ subnets_list = subnets.split("|") return FargateAgent( region_name=aws_region, cpu=agent_cpu, memory=agent_memory, cluster=cluster_name, taskRoleArn=task_role_arn, executionRoleArn=execution_role_arn, networkConfiguration={ "awsvpcConfiguration": { "assignPublicIp": "ENABLED", "subnets": subnets_list, "securityGroups": [], } }, containerDefinitions=[{ "logConfiguration": { "logDriver": "awslogs", "options": { "awslogs-region": aws_region, "awslogs-group": f"{environment}_dataflow_automation_agent", "awslogs-stream-prefix": "workflow_start", }, }, }], labels=[f"{environment}_dataflow_automation"], )
def test_deploy_flows_register_task_definition_all_args( monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.run_task.return_value = {} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent( aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", task_role_arn="task_role_arn", execution_role_arn="execution_role_arn", region_name="region", cluster="cluster", subnets=["subnet"], security_groups=["security_group"], repository_credentials="repo", assign_public_ip="DISABLED", task_cpu="1", task_memory="2", ) agent.deploy_flows(flow_runs=[ GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", }) ]) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert (boto3_client.register_task_definition.call_args[1]["family"] == "prefect-task-id") assert boto3_client.register_task_definition.call_args[1][ "containerDefinitions"] == [{ "name": "flow", "image": "test/name:tag", "command": ["/bin/sh", "-c", "prefect execute cloud-flow"], "environment": [ { "name": "PREFECT__CLOUD__API", "value": "https://api.prefect.io" }, { "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false" }, { "name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true" }, { "name": "PREFECT__LOGGING__LEVEL", "value": "DEBUG" }, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, ], "essential": True, "repositoryCredentials": { "credentialsParameter": "repo" }, }] assert boto3_client.register_task_definition.call_args[1][ "requiresCompatibilities"] == ["FARGATE"] assert boto3_client.register_task_definition.call_args[1][ "networkMode"] == "awsvpc" assert boto3_client.register_task_definition.call_args[1]["cpu"] == "1" assert boto3_client.register_task_definition.call_args[1]["memory"] == "2"
def test_deploy_flow_all_args(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {} boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) kwarg_dict = { "taskRoleArn": "test", "executionRoleArn": "test", "volumes": "test", "placementConstraints": "test", "cpu": "test", "memory": "test", "tags": "test", "pidMode": "test", "ipcMode": "test", "proxyConfiguration": "test", "inferenceAccelerators": "test", "cluster": "cluster", "count": "test", "startedBy": "test", "group": "test", "placementStrategy": "test", "platformVersion": "test", "networkConfiguration": { "awsvpcConfiguration": { "subnets": ["subnet"], "assignPublicIp": "DISABLED", "securityGroups": ["security_group"], } }, "enableECSManagedTags": "test", "propagateTags": "test", } agent = FargateAgent(aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", region_name="region", **kwarg_dict) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", })) assert boto3_client.describe_task_definition.called assert boto3_client.run_task.called assert boto3_client.run_task.call_args[1]["cluster"] == "cluster" assert boto3_client.run_task.call_args[1][ "taskDefinition"] == "prefect-task-id" assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE" assert boto3_client.run_task.call_args[1]["overrides"] == { "containerOverrides": [{ "name": "flow", "environment": [ { "name": "PREFECT__CLOUD__AUTH_TOKEN", "value": "" }, { "name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id" }, ], }] } assert boto3_client.run_task.call_args[1]["networkConfiguration"] == { "awsvpcConfiguration": { "subnets": ["subnet"], "assignPublicIp": "DISABLED", "securityGroups": ["security_group"], } }
def test_deploy_flows_all_args(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {} boto3_client.run_task.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent( aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", task_role_arn="task_role_arn", execution_role_arn="execution_role_arn", region_name="region", cluster="cluster", subnets=["subnet"], security_groups=["security_group"], repository_credentials="repo", assign_public_ip="DISABLED", task_cpu="1", task_memory="2", ) agent.deploy_flows(flow_runs=[ GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", }) ]) assert boto3_client.describe_task_definition.called assert boto3_client.run_task.called assert boto3_client.run_task.call_args[1]["cluster"] == "cluster" assert boto3_client.run_task.call_args[1][ "taskDefinition"] == "prefect-task-id" assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE" assert boto3_client.run_task.call_args[1]["overrides"] == { "containerOverrides": [{ "name": "flow", "environment": [ { "name": "PREFECT__CLOUD__AUTH_TOKEN", "value": "" }, { "name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id" }, ], }] } assert boto3_client.run_task.call_args[1]["networkConfiguration"] == { "awsvpcConfiguration": { "subnets": ["subnet"], "assignPublicIp": "DISABLED", "securityGroups": ["security_group"], } }
def test_deploy_flows_includes_agent_labels_in_environment( monkeypatch, runner_token, flag): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) kwarg_dict = { "taskRoleArn": "test", "executionRoleArn": "test", "volumes": "test", "placementConstraints": "test", "cpu": "1", "memory": "2", "tags": "test", "pidMode": "test", "ipcMode": "test", "proxyConfiguration": "test", "inferenceAccelerators": "test", "cluster": "cluster", "count": "test", "startedBy": "test", "group": "test", "placementStrategy": "test", "platformVersion": "test", "networkConfiguration": { "awsvpcConfiguration": { "subnets": ["subnet"], "assignPublicIp": "DISABLED", "securityGroups": ["security_group"], } }, "enableECSManagedTags": "test", "propagateTags": "test", } with set_temporary_config({"logging.log_to_cloud": flag}): agent = FargateAgent(aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="token", region_name="region", labels=["aws", "staging"], **kwarg_dict) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", }), "id": "id", })) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert (boto3_client.register_task_definition.call_args[1]["family"] == "prefect-task-id") assert boto3_client.register_task_definition.call_args[1][ "containerDefinitions"] == [{ "name": "flow", "image": "test/name:tag", "command": ["/bin/sh", "-c", "prefect execute cloud-flow"], "environment": [ { "name": "PREFECT__CLOUD__API", "value": "https://api.prefect.io" }, { "name": "PREFECT__CLOUD__AGENT__LABELS", "value": "['aws', 'staging']", }, { "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false" }, { "name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": str(flag).lower() }, { "name": "PREFECT__LOGGING__LEVEL", "value": "DEBUG" }, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, ], "essential": True, }] assert boto3_client.register_task_definition.call_args[1][ "requiresCompatibilities"] == ["FARGATE"] assert boto3_client.register_task_definition.call_args[1][ "networkMode"] == "awsvpc" assert boto3_client.register_task_definition.call_args[1]["cpu"] == "1" assert boto3_client.register_task_definition.call_args[1]["memory"] == "2"
def test_deploy_flows_enable_task_revisions_no_tags(monkeypatch, runner_token): boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = {"tags": []} boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]} boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) agent = FargateAgent(enable_task_revisions=True) agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "id", "version": 2, "name": "name", }), "id": "id", "name": "name", })) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called boto3_client.register_task_definition.assert_called_with( containerDefinitions=[{ "name": "flow", "image": "test/name:tag", "command": ["/bin/sh", "-c", "prefect execute cloud-flow"], "environment": [ { "name": "PREFECT__CLOUD__API", "value": "https://api.prefect.io" }, { "name": "PREFECT__CLOUD__AGENT__LABELS", "value": "[]" }, { "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false" }, { "name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "false" }, { "name": "PREFECT__LOGGING__LEVEL", "value": "DEBUG" }, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, ], "essential": True, }], family="name", networkMode="awsvpc", requiresCompatibilities=["FARGATE"], tags=[ { "key": "PrefectFlowId", "value": "id" }, { "key": "PrefectFlowVersion", "value": "2" }, ], ) assert boto3_client.run_task.called assert agent.task_definition_name == "name"