async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x): if x == 2: raise prefect.engine.signals.FAIL("Don't run for 2") elif not isinstance(x, int): return -99 return x + 10 with prefect.Flow( "simple map", run_config=prefect.run_configs.LocalRun(), storage=Local(directory=tmpdir), ) as flow: flow.numbers = numbers() flow.add1 = add.map(flow.numbers) flow.add2 = add.map(flow.add1) flow.add2.trigger = prefect.triggers.all_finished add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x, y): return x + y with prefect.Flow( "simple map", run_config=prefect.run_configs.LocalRun(), storage=Local(directory=tmpdir), ) as flow: flow.numbers1 = numbers() flow.numbers2 = numbers() flow.add = add.map(flow.numbers1, flow.numbers2) add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ flow = prefect.Flow( "diamond", storage=prefect.environments.storage.Local(directory=tmpdir), environment=prefect.environments.LocalEnvironment(), ) flow.a = prefect.Task("a") flow.b = prefect.Task("b") flow.c = prefect.Task("c") flow.d = prefect.Task("d") flow.add_edge(flow.a, flow.b) flow.add_edge(flow.a, flow.c) flow.add_edge(flow.b, flow.d) flow.add_edge(flow.c, flow.d) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x): if x == 2: raise prefect.engine.signals.FAIL("Don't run for 2") return x + 10 with prefect.Flow( "simple map", environment=LocalEnvironment(), storage=Local(directory=tmpdir), ) as flow: flow.numbers = numbers() flow.add1 = add.map(flow.numbers) flow.add2 = add.map(flow.add1) add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def test_exceptions_are_logged_and_not_thrown(self, monkeypatch): # Mock httpx post_mock = CoroutineMock() post_mock.side_effect = Exception("Boom! I'm testing an exception!") httpx_mock = MagicMock() async def post(*args, **kwargs): post_mock(*args, **kwargs) httpx_mock.post = post monkeypatch.setattr( "prefect_server.utilities.sens_o_matic_events.sens_o_matic_httpx_client", httpx_mock, ) with set_temporary_config("env", "production"): try: row_id = str(uuid.uuid4()) table_name = "test" result = await emit_delete_event(row_id=row_id, table_name=table_name) except Exception as e: pytest.fail("Unexpected error") await sleep(0) assert post_mock.call_args[0][0] == "https://sens-o-matic.prefect.io/"
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ @prefect.task def numbers(): return [1, 2, 3] @prefect.task def add(x): return x + 1 @prefect.task def get_sum(x): return sum(x) with prefect.Flow( "simple reduce", environment=LocalEnvironment(), storage=Local(directory=tmpdir), ) as flow: flow.numbers = numbers() flow.add = add.map(flow.numbers) flow.sum = get_sum(flow.add) add_state_handler(flow) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True)) return flow
async def test_handles_hasura_connection_error_as_api_error( self, monkeypatch, exc): post = CoroutineMock(side_effect=exc) monkeypatch.setattr("prefect_server.utilities.http.httpx_client.post", post) with set_temporary_config("hasura.execute_retry_seconds", 3): with pytest.raises(APIError): await hasura_client.execute("query { hello }", variables=dict(x=1, y=dict(z=2)))
async def test_event_payload_is_formed_properly(self, sens_o_matic_httpx_mock): with set_temporary_config("env", "production"): row_id = str(uuid.uuid4()) table_name = "test" result = await emit_delete_event(row_id=row_id, table_name=table_name) await sleep(0) evaluate_delete_event_payload( table_name=table_name, row_id=row_id, post_mock=sens_o_matic_httpx_mock )
async def test_event_payload_does_not_fire_for_local_env( self, sens_o_matic_httpx_mock ): result = False with set_temporary_config("env", "local"): row_id = str(uuid.uuid4()) table_name = "test" result = await emit_delete_event(row_id=row_id, table_name=table_name) await sleep(0) assert result is None assert sens_o_matic_httpx_mock.called is False
async def test_sendgrid_treats_to_email_address_without_at_sign_as_name( self, monkeypatch): """ Sendgrid appears to parse `from` emails but not `to` emails """ hook = MagicMock() monkeypatch.setattr("sendgrid.SendGridAPIClient.send", hook) with set_temporary_config("sendgrid.send_email", True): utilities.email.send_templated_email( from_email="from-email-addr", to_emails=["to-email-addr"], template_id=None, ) mail = hook.call_args[0][0] assert len(mail.personalizations) == 1 assert mail.personalizations[0].tos == [{"name": "to-email-addr"}]
async def test_handle_connection_error(self, monkeypatch): post = CoroutineMock(side_effect=lambda *args, **kwargs: MagicMock( json=MagicMock(side_effect=lambda: dict( data=Box( query=kwargs["json"]["query"], variables=kwargs["json"]["variables"], ), errors=[Box(message="connection error")], )))) monkeypatch.setattr("httpx.post", post) start_time = pendulum.now("utc") with set_temporary_config("hasura.execute_retry_seconds", 3): with pytest.raises(ValueError, match="Unable to connect to postgres"): await hasura.execute("query { hello }", variables=dict(x=1, y=dict(z=2))) # confirm we waited while retrying, leaving a couple of seconds to be conservative assert pendulum.now("utc") > start_time.add(seconds=2)
async def test_send_email(self, monkeypatch): hook = MagicMock() monkeypatch.setattr("sendgrid.SendGridAPIClient.send", hook) body = "Don't Panic" with set_temporary_config("sendgrid.send_email", True): utilities.email.send_email( from_email="*****@*****.**", to_emails=["*****@*****.**"], subject="Testing", body=body, ) assert hook.called is True mail = hook.call_args[0][0] assert mail.from_email.email == "*****@*****.**" assert mail.personalizations[0].tos == [{"email": "*****@*****.**"}] assert mail.subject.subject == "Testing" assert mail.contents[0].content == body
async def test_send_email_payload(self, monkeypatch): hook = MagicMock() monkeypatch.setattr("sendgrid.SendGridAPIClient.send", hook) with set_temporary_config("sendgrid.send_email", True): utilities.email.send_templated_email( from_email="*****@*****.**", to_emails=["*****@*****.**"], template_id="test-template-id", template_data={"x": "y"}, ) assert hook.called is True mail = hook.call_args[0][0] assert mail.from_email.email == "*****@*****.**" assert len(mail.personalizations) == 1 assert mail.personalizations[0].tos == [{ "email": "*****@*****.**" }] assert mail.template_id.template_id == "test-template-id" assert mail.personalizations[0].dynamic_template_data == {"x": "y"}
async def flow(self, project_id, tmpdir): """ A simple diamond flow """ flow = prefect.Flow( "pause", storage=prefect.storage.Local(directory=tmpdir), run_config=prefect.run_configs.LocalRun(), ) flow.a = prefect.Task("a") flow.b = prefect.Task("b", trigger=prefect.triggers.manual_only) flow.c = prefect.Task("c") flow.add_edge(flow.a, flow.b) flow.add_edge(flow.b, flow.c) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True) ) return flow
async def flow(self, project_id, tmpdir): """ A diamond flow whose tasks always fail the first time """ class FailOnceTask(prefect.Task): def __init__(self, name): super().__init__( name=name, retry_delay=datetime.timedelta(seconds=0), max_retries=1 ) def run(self): if prefect.context.task_run_count <= 1: raise ValueError("Run me again!") flow = prefect.Flow( "diamond fail once", storage=prefect.environments.storage.Local(directory=tmpdir), environment=prefect.environments.LocalEnvironment(), ) flow.a = FailOnceTask("a") flow.b = FailOnceTask("b") flow.c = FailOnceTask("c") flow.d = FailOnceTask("d") flow.add_edge(flow.a, flow.b) flow.add_edge(flow.a, flow.c) flow.add_edge(flow.b, flow.d) flow.add_edge(flow.c, flow.d) with set_temporary_config(key="dev", value=True): flow.server_id = await api.flows.create_flow( project_id=project_id, serialized_flow=flow.serialize(build=True) ) return flow