def test_inspect_pipelines(self, cell): with cell: # Cell 1 pipeline_1 = beam.Pipeline(ir.InteractiveRunner()) pipeline_2 = beam.Pipeline(ir.InteractiveRunner()) with cell: # Cell 2 # Re-executes the line that created pipeline_1 causing the original # pipeline_1 no longer inspectable. pipeline_1 = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) expected_inspectable_pipelines = { pipeline_1: 'pipeline_1', pipeline_2: 'pipeline_2' } for inspectable_pipeline in inspector.inspect_pipelines().items(): self.assertTrue( inspectable_pipeline in expected_inspectable_pipelines.items()) pipeline_1_metadata = inspector.meta('pipeline_1', pipeline_1) pipeline_2_metadata = inspector.meta('pipeline_2', pipeline_2) expected_inspectables = { obfuscate(pipeline_2_metadata): { 'metadata': pipeline_2_metadata, 'value': pipeline_2 }, obfuscate(pipeline_1_metadata): { 'metadata': pipeline_1_metadata, 'value': pipeline_1 } } for inspectable in inspector.inspect().items(): self.assertTrue(inspectable in expected_inspectables.items())
def test_inspect(self, cell): with cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # Early watch the pipeline so that cell re-execution can be handled. ib.watch({'pipeline': pipeline}) # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(range(10)) with cell: # Cell 2 # Re-executes the line that created the pcoll causing the original # pcoll no longer inspectable. # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(range(10)) ib.watch(locals()) pipeline_metadata = inspector.meta('pipeline', pipeline) pcoll_metadata = inspector.meta('pcoll', pcoll) expected_inspectables = { obfuscate(pipeline_metadata): { 'metadata': pipeline_metadata, 'value': pipeline }, obfuscate(pcoll_metadata): { 'metadata': pcoll_metadata, 'value': pcoll } } for inspectable in inspector.inspect().items(): self.assertTrue(inspectable in expected_inspectables.items())
def test_get_val(self, cell): with cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=range-builtin-not-iterating pcoll = pipeline | 'Create' >> beam.Create(range(10)) with cell: # Cell 2 # Re-executes the line that created pipeline causing the original # pipeline become an anonymous pipeline that is still inspectable because # its pcoll is still inspectable. pipeline = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) ins = inspector.InteractiveEnvironmentInspector() _ = ins.list_inspectables() pipeline_identifier = obfuscate(inspector.meta('pipeline', pipeline)) self.assertIs(ins.get_val(pipeline_identifier), pipeline) pcoll_identifier = obfuscate(inspector.meta('pcoll', pcoll)) self.assertIs(ins.get_val(pcoll_identifier), pcoll) anonymous_pipeline_name = inspector.synthesize_pipeline_name( pcoll.pipeline) anonymous_pipeline_identifier = obfuscate( inspector.meta(anonymous_pipeline_name, pcoll.pipeline)) self.assertIs(ins.get_val(anonymous_pipeline_identifier), pcoll.pipeline)
def test_list_inspectables(self, cell): with cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=range-builtin-not-iterating pcoll_1 = pipeline | 'Create' >> beam.Create(range(10)) pcoll_2 = pcoll_1 | 'Square' >> beam.Map(lambda x: x * x) with cell: # Cell 2 # Re-executes the line that created pipeline causing the original # pipeline become an anonymous pipeline that is still inspectable because # its pcoll_1 and pcoll_2 are still inspectable. pipeline = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) anonymous_pipeline_name = inspector.synthesize_pipeline_name( pcoll_1.pipeline) anonymous_pipeline_metadata = inspector.meta(anonymous_pipeline_name, pcoll_1.pipeline) pipeline_metadata = inspector.meta('pipeline', pipeline) pcoll_1_metadata = inspector.meta('pcoll_1', pcoll_1) pcoll_2_metadata = inspector.meta('pcoll_2', pcoll_2) expected_inspectable_list = { obfuscate(pipeline_metadata): { 'metadata': pipeline_metadata, 'pcolls': {} }, obfuscate(anonymous_pipeline_metadata): { 'metadata': anonymous_pipeline_metadata, 'pcolls': { obfuscate(pcoll_1_metadata): pcoll_1_metadata, obfuscate(pcoll_2_metadata): pcoll_2_metadata } } } ins = inspector.InteractiveEnvironmentInspector() actual_listings = ins.list_inspectables() self.assertEqual(actual_listings, json.dumps(expected_inspectable_list))
def test_get_pcoll_data(self): pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=range-builtin-not-iterating pcoll = pipeline | 'Create' >> beam.Create(range(10)) counts = pcoll | beam.combiners.Count.PerElement() ib.watch(locals()) counts_identifier = obfuscate(inspector.meta('counts', counts)) ins = inspector.InteractiveEnvironmentInspector() _ = ins.list_inspectables() actual_counts_pcoll_data = ins.get_pcoll_data(counts_identifier) expected_counts_pcoll_data = ib.collect(counts).to_json(orient='table') self.assertEqual(actual_counts_pcoll_data, expected_counts_pcoll_data) actual_counts_with_window_info = ins.get_pcoll_data(counts_identifier, True) expected_counts_with_window_info = ib.collect(counts, True).to_json(orient='table') self.assertEqual( actual_counts_with_window_info, expected_counts_with_window_info)
def test_get_pcoll_data(self): pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(list(range(10))) counts = pcoll | beam.combiners.Count.PerElement() ib.watch(locals()) ie.current_env().track_user_pipelines() counts_identifier = obfuscate(inspector.meta('counts', counts)) ins = inspector.InteractiveEnvironmentInspector() _ = ins.list_inspectables() actual_counts_pcoll_data = ins.get_pcoll_data(counts_identifier) expected_counts_pcoll_data = ib.collect( counts, n=10).to_json(orient='table') self.assertEqual(actual_counts_pcoll_data, expected_counts_pcoll_data) actual_counts_with_window_info = ins.get_pcoll_data(counts_identifier, True) expected_counts_with_window_info = ib.collect( counts, include_window_info=True).to_json(orient='table') self.assertEqual( actual_counts_with_window_info, expected_counts_with_window_info)