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_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(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
def test_deploy_flows_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 = {} 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][ "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_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_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): 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)) 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", } 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_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__AGENT__LABELS", "value": "['aws', 'staging']", }, {"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, } ] 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_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)) 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_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"], } }