Python DataflowPipelineResult.wait_until_finish Examples

Example #1

File: dataflow_runner_test.py Project: gamars/beam

  def test_wait_until_finish(self, patched_time_sleep):
    values_enum = dataflow_api.Job.CurrentStateValueValuesEnum

    class MockDataflowRunner(object):

      def __init__(self, final_state):
        self.dataflow_client = mock.MagicMock()
        self.job = mock.MagicMock()
        self.job.currentState = values_enum.JOB_STATE_UNKNOWN

        def get_job_side_effect(*args, **kwargs):
          self.job.currentState = final_state
          return mock.DEFAULT

        self.dataflow_client.get_job = mock.MagicMock(
            return_value=self.job, side_effect=get_job_side_effect)
        self.dataflow_client.list_messages = mock.MagicMock(
            return_value=([], None))

    with self.assertRaisesRegexp(
        DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'):
      failed_runner = MockDataflowRunner(values_enum.JOB_STATE_FAILED)
      failed_result = DataflowPipelineResult(failed_runner.job, failed_runner)
      failed_result.wait_until_finish()

    succeeded_runner = MockDataflowRunner(values_enum.JOB_STATE_DONE)
    succeeded_result = DataflowPipelineResult(
        succeeded_runner.job, succeeded_runner)
    succeeded_result.wait_until_finish()

Example #2

  def test_wait_until_finish(self, patched_time_sleep):
    values_enum = dataflow_api.Job.CurrentStateValueValuesEnum

    class MockDataflowRunner(object):

      def __init__(self, final_state):
        self.dataflow_client = mock.MagicMock()
        self.job = mock.MagicMock()
        self.job.currentState = values_enum.JOB_STATE_UNKNOWN

        def get_job_side_effect(*args, **kwargs):
          self.job.currentState = final_state
          return mock.DEFAULT

        self.dataflow_client.get_job = mock.MagicMock(
            return_value=self.job, side_effect=get_job_side_effect)
        self.dataflow_client.list_messages = mock.MagicMock(
            return_value=([], None))

    with self.assertRaisesRegexp(
        DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'):
      failed_runner = MockDataflowRunner(values_enum.JOB_STATE_FAILED)
      failed_result = DataflowPipelineResult(failed_runner.job, failed_runner)
      failed_result.wait_until_finish()

    succeeded_runner = MockDataflowRunner(values_enum.JOB_STATE_DONE)
    succeeded_result = DataflowPipelineResult(
        succeeded_runner.job, succeeded_runner)
    succeeded_result.wait_until_finish()

Example #3

File: dataflow_runner_test.py Project: eljefe6a/incubator-beam

 def _duration_failed():
   with self.assertRaisesRegexp(
       DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'):
     duration_failed_runner = MockDataflowRunner(
         [values_enum.JOB_STATE_FAILED])
     duration_failed_result = DataflowPipelineResult(
         duration_failed_runner.job, duration_failed_runner)
     duration_failed_result.wait_until_finish(5)

Example #4

File: dataflow_runner_test.py Project: lindaJang/beam

 def _duration_failed():
     with self.assertRaisesRegexp(
             DataflowRuntimeException,
             'Dataflow pipeline failed. State: FAILED'):
         duration_failed_runner = MockDataflowRunner(
             [values_enum.JOB_STATE_FAILED])
         duration_failed_result = DataflowPipelineResult(
             duration_failed_runner.job, duration_failed_runner)
         duration_failed_result.wait_until_finish(5)

Example #5

File: dataflow_runner_test.py Project: lindaJang/beam

 def _duration_timedout():
     duration_timedout_runner = MockDataflowRunner(
         [values_enum.JOB_STATE_RUNNING])
     duration_timedout_result = DataflowPipelineResult(
         duration_timedout_runner.job, duration_timedout_runner)
     result = duration_timedout_result.wait_until_finish(5)
     self.assertEqual(result, PipelineState.RUNNING)

Example #6

File: dataflow_runner_test.py Project: lindaJang/beam

 def _duration_succeeded():
     duration_succeeded_runner = MockDataflowRunner(
         [values_enum.JOB_STATE_RUNNING, values_enum.JOB_STATE_DONE])
     duration_succeeded_result = DataflowPipelineResult(
         duration_succeeded_runner.job, duration_succeeded_runner)
     result = duration_succeeded_result.wait_until_finish(5)
     self.assertEqual(result, PipelineState.DONE)

Example #7

File: dataflow_runner_test.py Project: eljefe6a/incubator-beam

 def _duration_timedout():
   duration_timedout_runner = MockDataflowRunner(
       [values_enum.JOB_STATE_RUNNING])
   duration_timedout_result = DataflowPipelineResult(
       duration_timedout_runner.job, duration_timedout_runner)
   result = duration_timedout_result.wait_until_finish(5)
   self.assertEqual(result, PipelineState.RUNNING)

Example #8

File: dataflow_runner_test.py Project: eljefe6a/incubator-beam

 def _duration_succeeded():
   duration_succeeded_runner = MockDataflowRunner(
       [values_enum.JOB_STATE_RUNNING, values_enum.JOB_STATE_DONE])
   duration_succeeded_result = DataflowPipelineResult(
       duration_succeeded_runner.job, duration_succeeded_runner)
   result = duration_succeeded_result.wait_until_finish(5)
   self.assertEqual(result, PipelineState.DONE)

Example #9

File: dataflow_runner_test.py Project: NarasimhaKattunga/kafka-1

  def test_wait_until_finish(self, patched_time_sleep):
    values_enum = dataflow_api.Job.CurrentStateValueValuesEnum

    class MockDataflowRunner(object):
      def __init__(self, states):
        self.dataflow_client = mock.MagicMock()
        self.job = mock.MagicMock()
        self.job.currentState = values_enum.JOB_STATE_UNKNOWN
        self._states = states
        self._next_state_index = 0

        def get_job_side_effect(*args, **kwargs):
          self.job.currentState = self._states[self._next_state_index]
          if self._next_state_index < (len(self._states) - 1):
            self._next_state_index += 1
          return mock.DEFAULT

        self.dataflow_client.get_job = mock.MagicMock(
            return_value=self.job, side_effect=get_job_side_effect)
        self.dataflow_client.list_messages = mock.MagicMock(
            return_value=([], None))

    with self.assertRaisesRegex(DataflowRuntimeException,
                                'Dataflow pipeline failed. State: FAILED'):
      failed_runner = MockDataflowRunner([values_enum.JOB_STATE_FAILED])
      failed_result = DataflowPipelineResult(failed_runner.job, failed_runner)
      failed_result.wait_until_finish()

    succeeded_runner = MockDataflowRunner([values_enum.JOB_STATE_DONE])
    succeeded_result = DataflowPipelineResult(
        succeeded_runner.job, succeeded_runner)
    result = succeeded_result.wait_until_finish()
    self.assertEqual(result, PipelineState.DONE)

    # Time array has duplicate items, because some logging implementations also
    # call time.
    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])):
      duration_succeeded_runner = MockDataflowRunner(
          [values_enum.JOB_STATE_RUNNING, values_enum.JOB_STATE_DONE])
      duration_succeeded_result = DataflowPipelineResult(
          duration_succeeded_runner.job, duration_succeeded_runner)
      result = duration_succeeded_result.wait_until_finish(5000)
      self.assertEqual(result, PipelineState.DONE)

    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20, 20])):
      duration_timedout_runner = MockDataflowRunner(
          [values_enum.JOB_STATE_RUNNING])
      duration_timedout_result = DataflowPipelineResult(
          duration_timedout_runner.job, duration_timedout_runner)
      result = duration_timedout_result.wait_until_finish(5000)
      self.assertEqual(result, PipelineState.RUNNING)

    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])):
      with self.assertRaisesRegex(DataflowRuntimeException,
                                  'Dataflow pipeline failed. State: CANCELLED'):
        duration_failed_runner = MockDataflowRunner(
            [values_enum.JOB_STATE_CANCELLED])
        duration_failed_result = DataflowPipelineResult(
            duration_failed_runner.job, duration_failed_runner)
        duration_failed_result.wait_until_finish(5000)

Example #10

File: dataflow_runner_test.py Project: aaltay/incubator-beam

  def test_wait_until_finish(self, patched_time_sleep):
    values_enum = dataflow_api.Job.CurrentStateValueValuesEnum

    class MockDataflowRunner(object):

      def __init__(self, states):
        self.dataflow_client = mock.MagicMock()
        self.job = mock.MagicMock()
        self.job.currentState = values_enum.JOB_STATE_UNKNOWN
        self._states = states
        self._next_state_index = 0

        def get_job_side_effect(*args, **kwargs):
          self.job.currentState = self._states[self._next_state_index]
          if self._next_state_index < (len(self._states) - 1):
            self._next_state_index += 1
          return mock.DEFAULT

        self.dataflow_client.get_job = mock.MagicMock(
            return_value=self.job, side_effect=get_job_side_effect)
        self.dataflow_client.list_messages = mock.MagicMock(
            return_value=([], None))

    with self.assertRaisesRegexp(
        DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'):
      failed_runner = MockDataflowRunner([values_enum.JOB_STATE_FAILED])
      failed_result = DataflowPipelineResult(failed_runner.job, failed_runner)
      failed_result.wait_until_finish()

    succeeded_runner = MockDataflowRunner([values_enum.JOB_STATE_DONE])
    succeeded_result = DataflowPipelineResult(
        succeeded_runner.job, succeeded_runner)
    result = succeeded_result.wait_until_finish()
    self.assertEqual(result, PipelineState.DONE)

    # Time array has duplicate items, because some logging implementations also
    # call time.
    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])):
      duration_succeeded_runner = MockDataflowRunner(
          [values_enum.JOB_STATE_RUNNING, values_enum.JOB_STATE_DONE])
      duration_succeeded_result = DataflowPipelineResult(
          duration_succeeded_runner.job, duration_succeeded_runner)
      result = duration_succeeded_result.wait_until_finish(5000)
      self.assertEqual(result, PipelineState.DONE)

    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20, 20])):
      duration_timedout_runner = MockDataflowRunner(
          [values_enum.JOB_STATE_RUNNING])
      duration_timedout_result = DataflowPipelineResult(
          duration_timedout_runner.job, duration_timedout_runner)
      result = duration_timedout_result.wait_until_finish(5000)
      self.assertEqual(result, PipelineState.RUNNING)

    with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])):
      with self.assertRaisesRegexp(
          DataflowRuntimeException,
          'Dataflow pipeline failed. State: CANCELLED'):
        duration_failed_runner = MockDataflowRunner(
            [values_enum.JOB_STATE_CANCELLED])
        duration_failed_result = DataflowPipelineResult(
            duration_failed_runner.job, duration_failed_runner)
        duration_failed_result.wait_until_finish(5000)