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))
Beispiel #2
0
  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)