def test_execute_run_task_agent_token(monkeypatch):
boto3_client = MagicMock()
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
with set_temporary_config({"cloud.agent.auth_token": "test"}):
environment = FargateTaskEnvironment(cluster="test",
family="test",
taskDefinition="test")
environment.execute(
storage=Docker(registry_url="test",
image_name="image",
image_tag="tag"),
flow_location=".prefect/flows",
)
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["taskDefinition"] == "test"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow-container",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": prefect.config.cloud.agent.get("auth_token"),
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "unknown"
},
{
"name": "PREFECT__CONTEXT__IMAGE",
"value": "test/image:tag"
},
{
"name": "PREFECT__CONTEXT__FLOW_FILE_PATH",
"value": ".prefect/flows",
},
],
}]
}
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["cluster"] == "test"
def test_entire_environment_process_together(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.register_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "session")
monkeypatch.setenv("REGION_NAME", "region")
with prefect.context({"flow_run_id": "id"}), set_temporary_config(
{"cloud.auth_token": "test", "logging.extra_loggers": "['test_logger']",}
):
storage = Docker(registry_url="test", image_name="image", image_tag="tag")
flow = Flow("name", storage=storage)
environment = FargateTaskEnvironment(
containerDefinitions=[
{
"name": "flow-container",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}
],
cluster="test",
family="test",
taskDefinition="test",
)
assert environment
assert environment.aws_access_key_id == "id"
assert environment.aws_secret_access_key == "secret"
assert environment.aws_session_token == "session"
assert environment.region_name == "region"
environment.setup(flow=flow)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1]["family"] == "test"
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"
] == [
{
"name": "flow-container",
"image": "test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.FargateTaskEnvironment().run_flow()'",
],
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": prefect.config.cloud.graphql,
},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": "['test_logger']",
},
],
"essential": True,
}
]
environment.execute(flow=flow)
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["taskDefinition"] == "test"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [
{
"name": "flow-container",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": prefect.config.cloud.get("auth_token"),
},
{"name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id"},
{"name": "PREFECT__CONTEXT__IMAGE", "value": "test/image:tag"},
],
}
]
}
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["cluster"] == "test"
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(prefect.Flow("name"))
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
flow_run=[
GraphQLResult(
{
"flow": GraphQLResult(
{"name": "name", "storage": d.serialize(),}
)
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
def test_entire_environment_process_together(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.register_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "session")
monkeypatch.setenv("REGION_NAME", "region")
with prefect.context({"flow_run_id": "id"}):
storage = Docker(registry_url="test",
image_name="image",
image_tag="tag")
environment = FargateTaskEnvironment(
containerDefinitions=[{
"name": "flow-container",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}],
cluster="test",
family="test",
taskDefinition="test",
)
assert environment
assert environment.aws_access_key_id == "id"
assert environment.aws_secret_access_key == "secret"
assert environment.aws_session_token == "session"
assert environment.region_name == "region"
environment.setup(storage=storage)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1][
"family"] == "test"
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"] == [{
"name":
"flow-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.Flow.load(prefect.context.flow_file_path).environment.run_flow()'",
],
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": prefect.config.cloud.graphql,
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
],
"essential":
True,
}]
environment.execute(storage=storage, flow_location=".prefect/flows")
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["taskDefinition"] == "test"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow-container",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": prefect.config.cloud.agent.auth_token,
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
{
"name": "PREFECT__CONTEXT__IMAGE",
"value": "test/image:tag"
},
{
"name": "PREFECT__CONTEXT__FLOW_FILE_PATH",
"value": ".prefect/flows",
},
],
}]
}
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["cluster"] == "test"
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "flow_env.prefect"), "w+"):
flow = prefect.Flow("test")
flow_path = os.path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
with set_temporary_config({"cloud.auth_token": "test"}):
with prefect.context(flow_file_path=os.path.join(
directory, "flow_env.prefect")):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "test"
