def build_graphql_query():
log_message_event_fragment = get_log_message_event_fragment()
step_event_fragment = get_step_event_fragment()
return '\n'.join((
QUERY_TEMPLATE.format(
step_event_fragment=step_event_fragment.include_key,
log_message_event_fragment=log_message_event_fragment.include_key,
),
step_event_fragment.fragment,
log_message_event_fragment.fragment,
))
def test_all_step_events(): # pylint: disable=too-many-locals
handle = ExecutionTargetHandle.for_pipeline_fn(define_test_events_pipeline)
pipeline = handle.build_pipeline_definition()
mode = pipeline.get_default_mode_name()
execution_plan = create_execution_plan(pipeline, {}, mode=mode)
step_levels = execution_plan.topological_step_levels()
run_config = RunConfig(
executor_config=InProcessExecutorConfig(raise_on_error=False),
storage_mode=RunStorageMode.FILESYSTEM,
)
unhandled_events = STEP_EVENTS.copy()
# Exclude types that are not step events
ignored_events = {
'LogMessageEvent',
'PipelineStartEvent',
'PipelineSuccessEvent',
'PipelineInitFailureEvent',
'PipelineFailureEvent',
}
step_event_fragment = get_step_event_fragment()
log_message_event_fragment = get_log_message_event_fragment()
query = '\n'.join(
(
PIPELINE_EXECUTION_QUERY_TEMPLATE.format(
step_event_fragment=step_event_fragment.include_key,
log_message_event_fragment=log_message_event_fragment.include_key,
),
step_event_fragment.fragment,
log_message_event_fragment.fragment,
)
)
event_counts = defaultdict(int)
for step_level in step_levels:
for step in step_level:
variables = {
'executionParams': {
'selector': {'name': pipeline.name},
'environmentConfigData': {'storage': {'filesystem': {}}},
'mode': mode,
'executionMetadata': {'runId': run_config.run_id},
'stepKeys': [step.key],
}
}
pipeline_run_storage = PipelineRunStorage()
res = execute_query(handle, query, variables, pipeline_run_storage=pipeline_run_storage)
# go through the same dict, decrement all the event records we've seen from the GraphQL
# response
if not res.get('errors'):
run_logs = res['data']['startPipelineExecution']['run']['logs']['nodes']
events = [
dagster_event_from_dict(e, pipeline.name)
for e in run_logs
if e['__typename'] not in ignored_events
]
for event in events:
key = event.step_key + '.' + event.event_type_value
event_counts[key] -= 1
unhandled_events -= {DagsterEventType(e.event_type_value) for e in events}
# build up a dict, incrementing all the event records we've produced in the run storage
logs = pipeline_run_storage.get_run_by_id(run_config.run_id).all_logs()
for log in logs:
if not log.dagster_event or (
DagsterEventType(log.dagster_event.event_type_value) not in STEP_EVENTS
):
continue
key = log.dagster_event.step_key + '.' + log.dagster_event.event_type_value
event_counts[key] += 1
# Ensure we've processed all the events that were generated in the run storage
assert sum(event_counts.values()) == 0
# Ensure we've handled the universe of event types
assert not unhandled_events