interactive_beam_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Tests for apache_beam.runners.interactive.interactive_beam."""
# pytype: skip-file

import dataclasses
import importlib
import sys
import time
import unittest
from concurrent.futures import TimeoutError
from typing import NamedTuple
from unittest.mock import ANY
from unittest.mock import MagicMock
from unittest.mock import call
from unittest.mock import patch

import apache_beam as beam
from apache_beam import dataframe as frames
from apache_beam.dataframe.frame_base import DeferredBase
from apache_beam.options.pipeline_options import FlinkRunnerOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.runners.interactive import interactive_beam as ib
from apache_beam.runners.interactive import interactive_environment as ie
from apache_beam.runners.interactive import interactive_runner as ir
from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager
from apache_beam.runners.interactive.dataproc.types import ClusterMetadata
from apache_beam.runners.interactive.options.capture_limiters import Limiter
from apache_beam.runners.interactive.recording_manager import AsyncComputationResult
from apache_beam.runners.interactive.testing.mock_env import isolated_env
from apache_beam.runners.runner import PipelineState
from apache_beam.testing.test_stream import TestStream


@dataclasses.dataclass
class MockClusterMetadata:
  master_url = 'mock_url'


class Record(NamedTuple):
  order_id: int
  product_id: int
  quantity: int


# The module name is also a variable in module.
_module_name = 'apache_beam.runners.interactive.interactive_beam_test'


def _get_watched_pcollections_with_variable_names():
  watched_pcollections = {}
  for watching in ie.current_env().watching():
    for key, val in watching:
      if hasattr(val, '__class__') and isinstance(val, beam.pvalue.PCollection):
        watched_pcollections[val] = key
  return watched_pcollections


@unittest.skipIf(
    not ie.current_env().is_interactive_ready,
    '[interactive] dependency is not installed.')
@isolated_env
class InteractiveBeamTest(unittest.TestCase):
  def setUp(self):
    self._var_in_class_instance = 'a var in class instance, not directly used'

  def tearDown(self):
    ib.options.capture_control.set_limiters_for_test([])

  def test_watch_main_by_default(self):
    test_env = ie.InteractiveEnvironment()
    # Current Interactive Beam env fetched and the test env are 2 instances.
    self.assertNotEqual(id(ie.current_env()), id(test_env))
    self.assertEqual(ie.current_env().watching(), test_env.watching())

  def test_watch_a_module_by_name(self):
    test_env = ie.InteractiveEnvironment()
    ib.watch(_module_name)
    test_env.watch(_module_name)
    self.assertEqual(ie.current_env().watching(), test_env.watching())

  def test_watch_a_module_by_module_object(self):
    test_env = ie.InteractiveEnvironment()
    module = importlib.import_module(_module_name)
    ib.watch(module)
    test_env.watch(module)
    self.assertEqual(ie.current_env().watching(), test_env.watching())

  def test_watch_locals(self):
    # test_env serves as local var too.
    test_env = ie.InteractiveEnvironment()
    ib.watch(locals())
    test_env.watch(locals())
    self.assertEqual(ie.current_env().watching(), test_env.watching())

  def test_watch_class_instance(self):
    test_env = ie.InteractiveEnvironment()
    ib.watch(self)
    test_env.watch(self)
    self.assertEqual(ie.current_env().watching(), test_env.watching())

  @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]")
  def test_show_always_watch_given_pcolls(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    # pylint: disable=bad-option-value
    pcoll = p | 'Create' >> beam.Create(range(10))
    # The pcoll is not watched since watch(locals()) is not explicitly called.
    self.assertFalse(pcoll in _get_watched_pcollections_with_variable_names())
    # The call of show watches pcoll.
    ib.watch({'p': p})
    ie.current_env().track_user_pipelines()
    ib.show(pcoll)
    self.assertTrue(pcoll in _get_watched_pcollections_with_variable_names())

  @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]")
  def test_show_mark_pcolls_computed_when_done(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    # pylint: disable=bad-option-value
    pcoll = p | 'Create' >> beam.Create(range(10))
    self.assertFalse(pcoll in ie.current_env().computed_pcollections)
    # The call of show marks pcoll computed.
    ib.watch(locals())
    ie.current_env().track_user_pipelines()
    ib.show(pcoll)
    self.assertTrue(pcoll in ie.current_env().computed_pcollections)

  @patch((
      'apache_beam.runners.interactive.interactive_beam.'
      'visualize_computed_pcoll'))
  def test_show_handles_dict_of_pcolls(self, mocked_visualize):
    p = beam.Pipeline(ir.InteractiveRunner())
    # pylint: disable=bad-option-value
    pcoll = p | 'Create' >> beam.Create(range(10))
    ib.watch(locals())
    ie.current_env().track_user_pipelines()
    ie.current_env().mark_pcollection_computed([pcoll])
    ie.current_env()._is_in_ipython = True
    ie.current_env()._is_in_notebook = True
    ib.show({'pcoll': pcoll})
    mocked_visualize.assert_called_once()

  @patch((
      'apache_beam.runners.interactive.interactive_beam.'
      'visualize_computed_pcoll'))
  def test_show_handles_iterable_of_pcolls(self, mocked_visualize):
    p = beam.Pipeline(ir.InteractiveRunner())
    # pylint: disable=bad-option-value
    pcoll = p | 'Create' >> beam.Create(range(10))
    ib.watch(locals())
    ie.current_env().track_user_pipelines()
    ie.current_env().mark_pcollection_computed([pcoll])
    ie.current_env()._is_in_ipython = True
    ie.current_env()._is_in_notebook = True
    ib.show([pcoll])
    mocked_visualize.assert_called_once()

  @patch('apache_beam.runners.interactive.interactive_beam.visualize')
  def test_show_handles_deferred_dataframes(self, mocked_visualize):
    p = beam.Pipeline(ir.InteractiveRunner())

    deferred = frames.convert.to_dataframe(p | beam.Create([Record(0, 0, 0)]))

    ib.watch(locals())
    ie.current_env().track_user_pipelines()
    ie.current_env()._is_in_ipython = True
    ie.current_env()._is_in_notebook = True
    ib.show(deferred)
    mocked_visualize.assert_called_once()

  @patch((
      'apache_beam.runners.interactive.interactive_beam.'
      'visualize_computed_pcoll'))
  def test_show_noop_when_pcoll_container_is_invalid(self, mocked_visualize):
    class SomeRandomClass:
      def __init__(self, pcoll):
        self._pcoll = pcoll

    p = beam.Pipeline(ir.InteractiveRunner())
    # pylint: disable=bad-option-value
    pcoll = p | 'Create' >> beam.Create(range(10))
    ie.current_env().mark_pcollection_computed([pcoll])
    ie.current_env()._is_in_ipython = True
    ie.current_env()._is_in_notebook = True
    self.assertRaises(ValueError, ib.show, SomeRandomClass(pcoll))
    mocked_visualize.assert_not_called()

  def test_recordings_describe(self):
    """Tests that getting the description works."""

    # Create the pipelines to test.
    p1 = beam.Pipeline(ir.InteractiveRunner())
    p2 = beam.Pipeline(ir.InteractiveRunner())

    ib.watch(locals())

    # Get the descriptions. This test is simple as there isn't much logic in the
    # method.
    self.assertEqual(ib.recordings.describe(p1)['size'], 0)
    self.assertEqual(ib.recordings.describe(p2)['size'], 0)

    all_descriptions = ib.recordings.describe()
    self.assertEqual(all_descriptions[p1]['size'], 0)
    self.assertEqual(all_descriptions[p2]['size'], 0)

    # Ensure that the variable name for the pipeline is set correctly.
    self.assertEqual(all_descriptions[p1]['pipeline_var'], 'p1')
    self.assertEqual(all_descriptions[p2]['pipeline_var'], 'p2')

  def test_recordings_clear(self):
    """Tests that clearing the pipeline is correctly forwarded."""

    # Create a basic pipeline to store something in the cache.
    p = beam.Pipeline(ir.InteractiveRunner())
    elem = p | beam.Create([1])
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    # This records the pipeline so that the cache size is > 0.
    ib.collect(elem)
    self.assertGreater(ib.recordings.describe(p)['size'], 0)

    # After clearing, the cache should be empty.
    ib.recordings.clear(p)
    self.assertEqual(ib.recordings.describe(p)['size'], 0)

  def test_recordings_record(self):
    """Tests that recording pipeline succeeds."""

    # Add the TestStream so that it can be cached.
    ib.options.recordable_sources.add(TestStream)

    # Create a pipeline with an arbitrary amonunt of elements.
    p = beam.Pipeline(
        ir.InteractiveRunner(), options=PipelineOptions(streaming=True))
    # pylint: disable=unused-variable
    _ = (p
         | TestStream()
             .advance_watermark_to(0)
             .advance_processing_time(1)
             .add_elements(list(range(10)))
             .advance_processing_time(1))  # yapf: disable
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    # Assert that the pipeline starts in a good state.
    self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.STOPPED)
    self.assertEqual(ib.recordings.describe(p)['size'], 0)

    # Create a lmiter that stops the background caching job when something is
    # written to cache. This is used to make ensure that the pipeline is
    # functioning properly and that there are no data races with the test.
    class SizeLimiter(Limiter):
      def __init__(self, pipeline):
        self.pipeline = pipeline
        self.should_trigger = False

      def is_triggered(self):
        return (
            ib.recordings.describe(self.pipeline)['size'] > 0 and
            self.should_trigger)

    limiter = SizeLimiter(p)
    ib.options.capture_control.set_limiters_for_test([limiter])

    # Assert that a recording can be started only once.
    self.assertTrue(ib.recordings.record(p))
    self.assertFalse(ib.recordings.record(p))
    self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.RUNNING)

    # Wait for the pipeline to start and write something to cache.
    limiter.should_trigger = True
    for _ in range(60):
      if limiter.is_triggered():
        break
      time.sleep(1)
    self.assertTrue(
        limiter.is_triggered(),
        'Test timed out waiting for limiter to be triggered. This indicates '
        'that the BackgroundCachingJob did not cache anything.')

    # Assert that a recording can be stopped and can't be started again until
    # after the cache is cleared.
    ib.recordings.stop(p)
    self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.STOPPED)
    self.assertFalse(ib.recordings.record(p))
    ib.recordings.clear(p)
    self.assertTrue(ib.recordings.record(p))
    ib.recordings.stop(p)

  def test_collect_raw_records_true(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data = list(range(5))
    pcoll = p | 'Create' >> beam.Create(data)
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    result = ib.collect(pcoll, raw_records=True)
    self.assertIsInstance(result, list)
    self.assertEqual(result, data)

    result_n = ib.collect(pcoll, n=3, raw_records=True)
    self.assertIsInstance(result_n, list)
    self.assertEqual(result_n, data[:3])

  def test_collect_raw_records_false(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data = list(range(5))
    pcoll = p | 'Create' >> beam.Create(data)
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    result = ib.collect(pcoll)
    self.assertNotIsInstance(result, list)
    self.assertTrue(
        hasattr(result, 'columns'), "Result should have 'columns' attribute")
    self.assertTrue(
        hasattr(result, 'values'), "Result should have 'values' attribute")

    result_n = ib.collect(pcoll, n=3)
    self.assertNotIsInstance(result_n, list)
    self.assertTrue(
        hasattr(result_n, 'columns'),
        "Result (n=3) should have 'columns' attribute")
    self.assertTrue(
        hasattr(result_n, 'values'),
        "Result (n=3) should have 'values' attribute")

  def test_collect_raw_records_true_multiple_pcolls(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data1 = list(range(3))
    data2 = [x * x for x in range(3)]
    pcoll1 = p | 'Create1' >> beam.Create(data1)
    pcoll2 = p | 'Create2' >> beam.Create(data2)
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    result = ib.collect(pcoll1, pcoll2, raw_records=True)
    self.assertIsInstance(result, tuple)
    self.assertEqual(len(result), 2)
    self.assertIsInstance(result[0], list)
    self.assertEqual(result[0], data1)
    self.assertIsInstance(result[1], list)
    self.assertEqual(result[1], data2)

  def test_collect_raw_records_false_multiple_pcolls(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data1 = list(range(3))
    data2 = [x * x for x in range(3)]
    pcoll1 = p | 'Create1' >> beam.Create(data1)
    pcoll2 = p | 'Create2' >> beam.Create(data2)
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    result = ib.collect(pcoll1, pcoll2)
    self.assertIsInstance(result, tuple)
    self.assertEqual(len(result), 2)
    self.assertNotIsInstance(result[0], list)
    self.assertTrue(hasattr(result[0], 'columns'))
    self.assertNotIsInstance(result[1], list)
    self.assertTrue(hasattr(result[1], 'columns'))

  def test_collect_raw_records_true_force_tuple(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data = list(range(5))
    pcoll = p | 'Create' >> beam.Create(data)
    ib.watch(locals())
    ie.current_env().track_user_pipelines()

    result = ib.collect(pcoll, raw_records=True, force_tuple=True)
    self.assertIsInstance(result, tuple)
    self.assertEqual(len(result), 1)
    self.assertIsInstance(result[0], list)
    self.assertEqual(result[0], data)


@unittest.skipIf(
    not ie.current_env().is_interactive_ready,
    '[interactive] dependency is not installed.')
@isolated_env
class InteractiveBeamClustersTest(unittest.TestCase):
  def setUp(self):
    self.current_env.options.cache_root = 'gs://fake'
    self.clusters = self.current_env.clusters

  def tearDown(self):
    self.current_env.options.cache_root = None

  def test_cluster_metadata_pass_through_metadata(self):
    cid = ClusterMetadata(project_id='test-project')
    meta = self.clusters.cluster_metadata(cid)
    self.assertIs(meta, cid)

  def test_cluster_metadata_identifies_pipeline(self):
    cid = beam.Pipeline()
    known_meta = ClusterMetadata(project_id='test-project')
    dcm = DataprocClusterManager(known_meta)
    self.clusters.pipelines[cid] = dcm

    meta = self.clusters.cluster_metadata(cid)
    self.assertIs(meta, known_meta)

  def test_cluster_metadata_identifies_master_url(self):
    cid = 'test-url'
    known_meta = ClusterMetadata(project_id='test-project')
    _ = DataprocClusterManager(known_meta)
    self.clusters.master_urls[cid] = known_meta

    meta = self.clusters.cluster_metadata(cid)
    self.assertIs(meta, known_meta)

  def test_cluster_metadata_default_value(self):
    cid_none = None
    cid_unknown_p = beam.Pipeline()
    cid_unknown_master_url = 'test-url'
    default_meta = ClusterMetadata(project_id='test-project')
    self.clusters.set_default_cluster(default_meta)

    self.assertIs(default_meta, self.clusters.cluster_metadata(cid_none))
    self.assertIs(default_meta, self.clusters.cluster_metadata(cid_unknown_p))
    self.assertIs(
        default_meta, self.clusters.cluster_metadata(cid_unknown_master_url))

  def test_create_a_new_cluster(self):
    meta = ClusterMetadata(project_id='test-project')
    _ = self.clusters.create(meta)

    # Derived fields are populated.
    self.assertTrue(meta.master_url.startswith('test-url'))
    self.assertEqual(meta.dashboard, 'test-dashboard')
    # The cluster is known.
    self.assertIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertIn(meta.master_url, self.clusters.master_urls)
    # The default cluster is updated to the created cluster.
    self.assertIs(meta, self.clusters.default_cluster_metadata)

  def test_create_but_reuse_a_known_cluster(self):
    known_meta = ClusterMetadata(
        project_id='test-project', region='test-region')
    known_dcm = DataprocClusterManager(known_meta)
    known_meta.master_url = 'test-url'
    self.clusters.set_default_cluster(known_meta)
    self.clusters.dataproc_cluster_managers[known_meta] = known_dcm
    self.clusters.master_urls[known_meta.master_url] = known_meta

    # Use an equivalent meta as the identifier to create a cluster.
    cid_meta = ClusterMetadata(
        project_id=known_meta.project_id,
        region=known_meta.region,
        cluster_name=known_meta.cluster_name)
    dcm = self.clusters.create(cid_meta)
    # The known cluster manager is returned.
    self.assertIs(dcm, known_dcm)

    # Then use an equivalent master_url as the identifier.
    cid_master_url = known_meta.master_url
    dcm = self.clusters.create(cid_master_url)
    self.assertIs(dcm, known_dcm)

  def test_cleanup_by_a_pipeline(self):
    meta = ClusterMetadata(project_id='test-project')
    dcm = self.clusters.create(meta)

    # Set up the association between a pipeline and a cluster.
    # In real code, it's set by the runner the 1st time a pipeline is executed.
    options = PipelineOptions()
    options.view_as(FlinkRunnerOptions).flink_master = meta.master_url
    p = beam.Pipeline(options=options)
    self.clusters.pipelines[p] = dcm
    dcm.pipelines.add(p)

    self.clusters.cleanup(p)
    # Delete the cluster.
    self.m_delete_cluster.assert_called_once()
    # Pipeline association is cleaned up.
    self.assertNotIn(p, self.clusters.pipelines)
    self.assertNotIn(p, dcm.pipelines)
    # The internal option in the pipeline is overwritten.
    self.assertEqual(
        p.options.view_as(FlinkRunnerOptions).flink_master, '[auto]')
    # The original option is unchanged.
    self.assertEqual(
        options.view_as(FlinkRunnerOptions).flink_master, meta.master_url)
    # The cluster is unknown now.
    self.assertNotIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertNotIn(meta.master_url, self.clusters.master_urls)
    # The cleaned up cluster is also the default cluster. Clean the default.
    self.assertIsNone(self.clusters.default_cluster_metadata)

  def test_not_cleanup_if_multiple_pipelines_share_a_manager(self):
    meta = ClusterMetadata(project_id='test-project')
    dcm = self.clusters.create(meta)

    options = PipelineOptions()
    options.view_as(FlinkRunnerOptions).flink_master = meta.master_url
    options2 = PipelineOptions()
    options2.view_as(FlinkRunnerOptions).flink_master = meta.master_url
    p = beam.Pipeline(options=options)
    p2 = beam.Pipeline(options=options2)
    self.clusters.pipelines[p] = dcm
    self.clusters.pipelines[p2] = dcm
    dcm.pipelines.add(p)
    dcm.pipelines.add(p2)

    self.clusters.cleanup(p)
    # No cluster deleted.
    self.m_delete_cluster.assert_not_called()
    # Pipeline association of p is cleaned up.
    self.assertNotIn(p, self.clusters.pipelines)
    self.assertNotIn(p, dcm.pipelines)
    # The internal option in the pipeline is overwritten.
    self.assertEqual(
        p.options.view_as(FlinkRunnerOptions).flink_master, '[auto]')
    # The original option is unchanged.
    self.assertEqual(
        options.view_as(FlinkRunnerOptions).flink_master, meta.master_url)
    # Pipeline association of p2 still presents.
    self.assertIn(p2, self.clusters.pipelines)
    self.assertIn(p2, dcm.pipelines)
    self.assertEqual(
        p2.options.view_as(FlinkRunnerOptions).flink_master, meta.master_url)
    self.assertEqual(
        options2.view_as(FlinkRunnerOptions).flink_master, meta.master_url)
    # The cluster is still known.
    self.assertIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertIn(meta.master_url, self.clusters.master_urls)
    # The default cluster still presents.
    self.assertIs(meta, self.clusters.default_cluster_metadata)

  def test_cleanup_by_a_master_url(self):
    meta = ClusterMetadata(project_id='test-project')
    _ = self.clusters.create(meta)

    self.clusters.cleanup(meta.master_url)
    self.m_delete_cluster.assert_called_once()
    self.assertNotIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertNotIn(meta.master_url, self.clusters.master_urls)
    self.assertIsNone(self.clusters.default_cluster_metadata)

  def test_cleanup_by_meta(self):
    known_meta = ClusterMetadata(
        project_id='test-project', region='test-region')
    _ = self.clusters.create(known_meta)

    meta = ClusterMetadata(
        project_id=known_meta.project_id,
        region=known_meta.region,
        cluster_name=known_meta.cluster_name)
    self.clusters.cleanup(meta)
    self.m_delete_cluster.assert_called_once()
    self.assertNotIn(known_meta, self.clusters.dataproc_cluster_managers)
    self.assertNotIn(known_meta.master_url, self.clusters.master_urls)
    self.assertIsNone(self.clusters.default_cluster_metadata)

  def test_force_cleanup_everything(self):
    meta = ClusterMetadata(project_id='test-project')
    meta2 = ClusterMetadata(project_id='test-project-2')
    _ = self.clusters.create(meta)
    _ = self.clusters.create(meta2)

    self.clusters.cleanup(force=True)
    self.assertEqual(self.m_delete_cluster.call_count, 2)
    self.assertNotIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertNotIn(meta2, self.clusters.dataproc_cluster_managers)
    self.assertIsNone(self.clusters.default_cluster_metadata)

  def test_cleanup_noop_for_no_cluster_identifier(self):
    meta = ClusterMetadata(project_id='test-project')
    _ = self.clusters.create(meta)

    self.clusters.cleanup()
    self.m_delete_cluster.assert_not_called()

  def test_cleanup_noop_unknown_cluster(self):
    meta = ClusterMetadata(project_id='test-project')
    dcm = self.clusters.create(meta)
    p = beam.Pipeline()
    self.clusters.pipelines[p] = dcm
    dcm.pipelines.add(p)

    cid_pipeline = beam.Pipeline()
    self.clusters.cleanup(cid_pipeline)
    self.m_delete_cluster.assert_not_called()

    cid_master_url = 'some-random-url'
    self.clusters.cleanup(cid_master_url)
    self.m_delete_cluster.assert_not_called()

    cid_meta = ClusterMetadata(project_id='random-project')
    self.clusters.cleanup(cid_meta)
    self.m_delete_cluster.assert_not_called()

    self.assertIn(meta, self.clusters.dataproc_cluster_managers)
    self.assertIn(meta.master_url, self.clusters.master_urls)
    self.assertIs(meta, self.clusters.default_cluster_metadata)
    self.assertIn(p, self.clusters.pipelines)
    self.assertIn(p, dcm.pipelines)

  def test_describe_everything(self):
    meta = ClusterMetadata(project_id='test-project')
    meta2 = ClusterMetadata(
        project_id='test-project', region='some-other-region')
    _ = self.clusters.create(meta)
    _ = self.clusters.create(meta2)

    meta_list = self.clusters.describe()
    self.assertEqual([meta, meta2], meta_list)

  def test_describe_by_cluster_identifier(self):
    known_meta = ClusterMetadata(project_id='test-project')
    known_meta2 = ClusterMetadata(
        project_id='test-project', region='some-other-region')
    dcm = self.clusters.create(known_meta)
    dcm2 = self.clusters.create(known_meta2)
    p = beam.Pipeline()
    p2 = beam.Pipeline()
    self.clusters.pipelines[p] = dcm
    dcm.pipelines.add(p)
    self.clusters.pipelines[p2] = dcm2
    dcm.pipelines.add(p2)

    cid_pipeline = p
    meta = self.clusters.describe(cid_pipeline)
    self.assertIs(meta, known_meta)

    cid_master_url = known_meta.master_url
    meta = self.clusters.describe(cid_master_url)
    self.assertIs(meta, known_meta)

    cid_meta = ClusterMetadata(
        project_id=known_meta.project_id,
        region=known_meta.region,
        cluster_name=known_meta.cluster_name)
    meta = self.clusters.describe(cid_meta)
    self.assertIs(meta, known_meta)

  def test_describe_everything_when_cluster_identifer_unknown(self):
    known_meta = ClusterMetadata(project_id='test-project')
    known_meta2 = ClusterMetadata(
        project_id='test-project', region='some-other-region')
    dcm = self.clusters.create(known_meta)
    dcm2 = self.clusters.create(known_meta2)
    p = beam.Pipeline()
    p2 = beam.Pipeline()
    self.clusters.pipelines[p] = dcm
    dcm.pipelines.add(p)
    self.clusters.pipelines[p2] = dcm2
    dcm.pipelines.add(p2)

    cid_pipeline = beam.Pipeline()
    meta_list = self.clusters.describe(cid_pipeline)
    self.assertEqual([known_meta, known_meta2], meta_list)

    cid_master_url = 'some-random-url'
    meta_list = self.clusters.describe(cid_master_url)
    self.assertEqual([known_meta, known_meta2], meta_list)

    cid_meta = ClusterMetadata(project_id='some-random-project')
    meta_list = self.clusters.describe(cid_meta)
    self.assertEqual([known_meta, known_meta2], meta_list)

  def test_default_value_for_invalid_worker_number(self):
    meta = ClusterMetadata(project_id='test-project', num_workers=1)
    self.clusters.create(meta)

    self.assertEqual(meta.num_workers, 2)


@unittest.skipIf(
    not ie.current_env().is_interactive_ready,
    '[interactive] dependency is not installed.')
@isolated_env
class InteractiveBeamComputeTest(unittest.TestCase):
  def setUp(self):
    self.env = ie.current_env()
    self.env._is_in_ipython = False  # Default to non-IPython

  def test_compute_blocking(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data = list(range(10))
    pcoll = p | 'Create' >> beam.Create(data)
    ib.watch(locals())
    self.env.track_user_pipelines()

    result = ib.compute(pcoll, blocking=True)
    self.assertIsNone(result)  # Blocking returns None
    self.assertTrue(pcoll in self.env.computed_pcollections)
    collected = ib.collect(pcoll, raw_records=True)
    self.assertEqual(collected, data)

  def test_compute_non_blocking(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    data = list(range(5))
    pcoll = p | 'Create' >> beam.Create(data)
    ib.watch(locals())
    self.env.track_user_pipelines()

    async_result = ib.compute(pcoll, blocking=False)
    self.assertIsInstance(async_result, AsyncComputationResult)

    pipeline_result = async_result.result(timeout=60)
    self.assertTrue(async_result.done())
    self.assertIsNone(async_result.exception())
    self.assertEqual(pipeline_result.state, PipelineState.DONE)
    self.assertTrue(pcoll in self.env.computed_pcollections)
    collected = ib.collect(pcoll, raw_records=True)
    self.assertEqual(collected, data)

  def test_compute_with_list_input(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
    pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6])
    ib.watch(locals())
    self.env.track_user_pipelines()

    ib.compute([pcoll1, pcoll2], blocking=True)
    self.assertTrue(pcoll1 in self.env.computed_pcollections)
    self.assertTrue(pcoll2 in self.env.computed_pcollections)

  def test_compute_with_dict_input(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
    pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6])
    ib.watch(locals())
    self.env.track_user_pipelines()

    ib.compute({'a': pcoll1, 'b': pcoll2}, blocking=True)
    self.assertTrue(pcoll1 in self.env.computed_pcollections)
    self.assertTrue(pcoll2 in self.env.computed_pcollections)

  def test_compute_empty_input(self):
    result = ib.compute([], blocking=True)
    self.assertIsNone(result)
    result_async = ib.compute([], blocking=False)
    self.assertIsNone(result_async)

  def test_compute_force_recompute(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll = p | 'Create' >> beam.Create([1, 2, 3])
    ib.watch(locals())
    self.env.track_user_pipelines()

    ib.compute(pcoll, blocking=True)
    self.assertTrue(pcoll in self.env.computed_pcollections)

    # Mock evict_computed_pcollections to check if it's called
    with patch.object(self.env, 'evict_computed_pcollections') as mock_evict:
      ib.compute(pcoll, blocking=True, force_compute=True)
      mock_evict.assert_called_once_with(p)
    self.assertTrue(pcoll in self.env.computed_pcollections)

  def test_compute_non_blocking_exception(self):
    p = beam.Pipeline(ir.InteractiveRunner())

    def raise_error(elem):
      raise ValueError('Test Error')

    pcoll = p | 'Create' >> beam.Create([1]) | 'Error' >> beam.Map(raise_error)
    ib.watch(locals())
    self.env.track_user_pipelines()

    async_result = ib.compute(pcoll, blocking=False)
    self.assertIsInstance(async_result, AsyncComputationResult)

    with self.assertRaises(ValueError):
      async_result.result(timeout=60)

    self.assertTrue(async_result.done())
    self.assertIsInstance(async_result.exception(), ValueError)
    self.assertFalse(pcoll in self.env.computed_pcollections)

  @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True)
  @patch('apache_beam.runners.interactive.recording_manager.display')
  @patch('apache_beam.runners.interactive.recording_manager.clear_output')
  @patch('apache_beam.runners.interactive.recording_manager.HTML')
  @patch('ipywidgets.Button')
  @patch('ipywidgets.FloatProgress')
  @patch('ipywidgets.Output')
  @patch('ipywidgets.HBox')
  @patch('ipywidgets.VBox')
  def test_compute_non_blocking_ipython_widgets(
      self,
      mock_vbox,
      mock_hbox,
      mock_output,
      mock_progress,
      mock_button,
      mock_html,
      mock_clear_output,
      mock_display,
  ):
    self.env._is_in_ipython = True
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll = p | 'Create' >> beam.Create(range(3))
    ib.watch(locals())
    self.env.track_user_pipelines()

    mock_controls = mock_vbox.return_value
    mock_html_instance = mock_html.return_value

    async_result = ib.compute(pcoll, blocking=False)
    self.assertIsNotNone(async_result)
    mock_button.assert_called_once_with(description='Cancel')
    mock_progress.assert_called_once()
    mock_output.assert_called_once()
    mock_hbox.assert_called_once()
    mock_vbox.assert_called_once()
    mock_html.assert_called_once_with('<p>Initializing...</p>')

    self.assertEqual(mock_display.call_count, 2)
    mock_display.assert_has_calls([
        call(mock_controls, display_id=async_result._display_id),
        call(mock_html_instance)
    ])

    mock_clear_output.assert_called_once()
    async_result.wait_for_completion(timeout=60)

  def test_compute_dependency_wait_true(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
    pcoll2 = pcoll1 | 'Map' >> beam.Map(lambda x: x * 2)
    ib.watch(locals())
    self.env.track_user_pipelines()

    rm = self.env.get_recording_manager(p)

    # Start pcoll1 computation
    async_res1 = ib.compute(pcoll1, blocking=False)
    self.assertTrue(self.env.is_pcollection_computing(pcoll1))

    # Spy on _wait_for_dependencies
    with patch.object(rm,
                      '_wait_for_dependencies',
                      wraps=rm._wait_for_dependencies) as spy_wait:
      async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=True)

      # Check that wait_for_dependencies was called for pcoll2
      spy_wait.assert_called_with({pcoll2}, async_res2)

      # Let pcoll1 finish
      async_res1.wait_for_completion(timeout=60)
      self.assertTrue(pcoll1 in self.env.computed_pcollections)
      self.assertFalse(self.env.is_pcollection_computing(pcoll1))

      # pcoll2 should now run and complete
      async_res2.wait_for_completion(timeout=60)
      self.assertTrue(pcoll2 in self.env.computed_pcollections)

  @patch.object(ie.InteractiveEnvironment, 'is_pcollection_computing')
  def test_compute_dependency_wait_false(self, mock_is_computing):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
    pcoll2 = pcoll1 | 'Map' >> beam.Map(lambda x: x * 2)
    ib.watch(locals())
    self.env.track_user_pipelines()

    rm = self.env.get_recording_manager(p)

    # Pretend pcoll1 is computing
    mock_is_computing.side_effect = lambda pcoll: pcoll is pcoll1

    with patch.object(rm,
                      '_execute_pipeline_fragment',
                      wraps=rm._execute_pipeline_fragment) as spy_execute:
      async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=False)
      async_res2.wait_for_completion(timeout=60)

      # Assert that execute was called for pcoll2 without waiting
      spy_execute.assert_called_with({pcoll2}, async_res2, ANY, ANY)
      self.assertTrue(pcoll2 in self.env.computed_pcollections)

  def test_async_computation_result_cancel(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    # A stream that never finishes to test cancellation
    pcoll = p | beam.Create([1]) | beam.Map(lambda x: time.sleep(100))
    ib.watch(locals())
    self.env.track_user_pipelines()

    async_result = ib.compute(pcoll, blocking=False)
    self.assertIsInstance(async_result, AsyncComputationResult)

    # Give it a moment to start
    time.sleep(0.1)

    # Mock the pipeline result's cancel method
    mock_pipeline_result = MagicMock()
    mock_pipeline_result.state = PipelineState.RUNNING
    async_result.set_pipeline_result(mock_pipeline_result)

    self.assertTrue(async_result.cancel())
    mock_pipeline_result.cancel.assert_called_once()

    # The future should be cancelled eventually by the runner
    # This part is hard to test without deeper runner integration
    with self.assertRaises(TimeoutError):
      async_result.result(timeout=1)  # It should not complete successfully

  @patch(
      'apache_beam.runners.interactive.recording_manager.RecordingManager.'
      '_execute_pipeline_fragment')
  def test_compute_multiple_async(self, mock_execute_fragment):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3])
    pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6])
    pcoll3 = pcoll1 | 'Map1' >> beam.Map(lambda x: x * 2)
    ib.watch(locals())
    self.env.track_user_pipelines()

    mock_pipeline_result = MagicMock()
    mock_pipeline_result.state = PipelineState.DONE
    mock_execute_fragment.return_value = mock_pipeline_result

    res1 = ib.compute(pcoll1, blocking=False)
    res2 = ib.compute(pcoll2, blocking=False)
    res3 = ib.compute(pcoll3, blocking=False)  # Depends on pcoll1

    self.assertIsNotNone(res1)
    self.assertIsNotNone(res2)
    self.assertIsNotNone(res3)

    res1.result(timeout=60)
    res2.result(timeout=60)
    res3.result(timeout=60)

    time.sleep(0.1)

    self.assertTrue(
        pcoll1 in self.env.computed_pcollections, "pcoll1 not marked computed")
    self.assertTrue(
        pcoll2 in self.env.computed_pcollections, "pcoll2 not marked computed")
    self.assertTrue(
        pcoll3 in self.env.computed_pcollections, "pcoll3 not marked computed")

    self.assertEqual(mock_execute_fragment.call_count, 3)

  @patch(
      'apache_beam.runners.interactive.interactive_beam.'
      'deferred_df_to_pcollection')
  def test_compute_input_flattening(self, mock_deferred_to_pcoll):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p | 'C1' >> beam.Create([1])
    pcoll2 = p | 'C2' >> beam.Create([2])
    pcoll3 = p | 'C3' >> beam.Create([3])
    pcoll4 = p | 'C4' >> beam.Create([4])

    class MockDeferred(DeferredBase):
      def __init__(self, pcoll):
        mock_expr = MagicMock()
        super().__init__(mock_expr)
        self._pcoll = pcoll

      def _get_underlying_pcollection(self):
        return self._pcoll

    deferred_pcoll = MockDeferred(pcoll4)

    mock_deferred_to_pcoll.return_value = (pcoll4, p)

    ib.watch(locals())
    self.env.track_user_pipelines()

    with patch.object(self.env, 'get_recording_manager') as mock_get_rm:
      mock_rm = MagicMock()
      mock_get_rm.return_value = mock_rm
      ib.compute(pcoll1, [pcoll2], {'a': pcoll3}, deferred_pcoll)

      expected_pcolls = {pcoll1, pcoll2, pcoll3, pcoll4}
      mock_rm.compute_async.assert_called_once_with(
          expected_pcolls,
          wait_for_inputs=True,
          blocking=False,
          runner=None,
          options=None,
          force_compute=False)

  def test_compute_invalid_input_type(self):
    with self.assertRaisesRegex(ValueError,
                                "not a dict, an iterable or a PCollection"):
      ib.compute(123)

  def test_compute_mixed_pipelines(self):
    p1 = beam.Pipeline(ir.InteractiveRunner())
    pcoll1 = p1 | 'C1' >> beam.Create([1])
    p2 = beam.Pipeline(ir.InteractiveRunner())
    pcoll2 = p2 | 'C2' >> beam.Create([2])
    ib.watch(locals())
    self.env.track_user_pipelines()

    with self.assertRaisesRegex(
        ValueError, "All PCollections must belong to the same pipeline"):
      ib.compute(pcoll1, pcoll2)

  @patch(
      'apache_beam.runners.interactive.interactive_beam.'
      'deferred_df_to_pcollection')
  @patch.object(ib, 'watch')
  def test_compute_with_deferred_base(self, mock_watch, mock_deferred_to_pcoll):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll = p | 'C1' >> beam.Create([1])

    class MockDeferred(DeferredBase):
      def __init__(self, pcoll):
        # Provide a dummy expression to satisfy DeferredBase.__init__
        mock_expr = MagicMock()
        super().__init__(mock_expr)
        self._pcoll = pcoll

      def _get_underlying_pcollection(self):
        return self._pcoll

    deferred = MockDeferred(pcoll)

    mock_deferred_to_pcoll.return_value = (pcoll, p)

    with patch.object(self.env, 'get_recording_manager') as mock_get_rm:
      mock_rm = MagicMock()
      mock_get_rm.return_value = mock_rm
      ib.compute(deferred)

      mock_deferred_to_pcoll.assert_called_once_with(deferred)
      self.assertEqual(mock_watch.call_count, 2)
      mock_watch.assert_has_calls([
          call({f'anonymous_pcollection_{id(pcoll)}': pcoll}),
          call({f'anonymous_pipeline_{id(p)}': p})
      ],
                                  any_order=False)
      mock_rm.compute_async.assert_called_once_with({pcoll},
                                                    wait_for_inputs=True,
                                                    blocking=False,
                                                    runner=None,
                                                    options=None,
                                                    force_compute=False)

  def test_compute_new_pipeline(self):
    p = beam.Pipeline(ir.InteractiveRunner())
    pcoll = p | 'Create' >> beam.Create([1])
    # NOT calling ib.watch() or track_user_pipelines()

    with patch.object(self.env, 'get_recording_manager') as mock_get_rm, \
        patch.object(ib, 'watch') as mock_watch:
      mock_rm = MagicMock()
      mock_get_rm.return_value = mock_rm
      ib.compute(pcoll)

      mock_watch.assert_called_with({f'anonymous_pipeline_{id(p)}': p})
      mock_get_rm.assert_called_once_with(p, create_if_absent=True)
      mock_rm.compute_async.assert_called_once()


if __name__ == '__main__':
  unittest.main()