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_register_task_definition(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() 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")
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_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_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_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_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_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"
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_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_disable_task_revisions_with_external_kwargs( monkeypatch, runner_token): boto3_client = MagicMock() boto3_resource = MagicMock() streaming_body = MagicMock() streaming_body.read.return_value.decode.return_value = '{"cpu": "256", "networkConfiguration": "test", "tags": [{"key": "test", "value": "test"}]}' boto3_resource.return_value.Object.return_value.get.return_value = { "Body": streaming_body } boto3_client.describe_task_definition.return_value = {} 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)) monkeypatch.setattr("boto3.resource", boto3_resource) agent = FargateAgent( enable_task_revisions=False, 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", cluster="test", labels=["aws", "staging"], ) agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "new_id", "version": 6, "name": "name", }), "id": "id", "name": "name", })) assert agent.task_definition_kwargs == {} assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.not_called boto3_client.run_task.assert_called_with( launchType="FARGATE", networkConfiguration="test", cluster="test", overrides={ "containerOverrides": [{ "name": "flow", "environment": [ { "name": "PREFECT__CLOUD__AUTH_TOKEN", "value": "" }, { "name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id" }, ], }] }, taskDefinition="prefect-task-new_id", tags=[{ "key": "test", "value": "test" }], ) assert agent.task_definition_name == "prefect-task-new_id"
def test_deploy_flows_enable_task_revisions_with_external_kwargs( monkeypatch, runner_token): boto3_client = MagicMock() boto3_resource = MagicMock() streaming_body = MagicMock() streaming_body.read.return_value.decode.return_value = '{"cpu": "256", "networkConfiguration": "test", "tags": [{"key": "test", "value": "test"}]}' boto3_resource.return_value.Object.return_value.get.return_value = { "Body": streaming_body } boto3_client.describe_task_definition.return_value = { "tags": [ { "key": "PrefectFlowId", "value": "id" }, { "key": "PrefectFlowVersion", "value": "5" }, ] } 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)) monkeypatch.setattr("boto3.resource", boto3_resource) agent = FargateAgent( enable_task_revisions=True, 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", cluster="test", tags=[{ "key": "team", "value": "data" }], labels=["aws", "staging"], ) agent.deploy_flow( GraphQLResult({ "flow": GraphQLResult({ "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "id": "new_id", "version": 6, "name": "name", }), "id": "id", "name": "name", })) assert boto3_client.describe_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": "['aws', 'staging']", }, { "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, }], cpu="256", family="name", networkMode="awsvpc", requiresCompatibilities=["FARGATE"], tags=[ { "key": "test", "value": "test" }, { "key": "PrefectFlowId", "value": "new_id" }, { "key": "PrefectFlowVersion", "value": "6" }, ], ) boto3_client.run_task.assert_called_with( launchType="FARGATE", networkConfiguration="test", cluster="test", overrides={ "containerOverrides": [{ "name": "flow", "environment": [ { "name": "PREFECT__CLOUD__AUTH_TOKEN", "value": "" }, { "name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id" }, ], }] }, taskDefinition="name", tags=[{ "key": "test", "value": "test" }], ) assert agent.task_definition_name == "name"