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