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"