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_val(self, cell):
with cell: # Cell 1
pipeline = beam.Pipeline(ir.InteractiveRunner())
# pylint: disable=bad-option-value
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_get_pcoll_data(self):
pipeline = beam.Pipeline(ir.InteractiveRunner())
# pylint: disable=range-builtin-not-iterating
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)