def main():
properties, unused_args = parse_properties(sys.argv[1:])
# Initialize the logging machinery.
job_id = properties['job_id']
worker_id = properties['worker_id']
log_path = properties['dataflow.worker.logging.location']
logger.initialize(job_id, worker_id, log_path)
logging.info('Worker started with properties: %s', properties)
sdk_pipeline_options = json.loads(
properties.get('sdk_pipeline_options', '{}'))
logging.info('Worker started with sdk_pipeline_options: %s',
sdk_pipeline_options)
if unused_args:
logging.warning('Unrecognized arguments %s', unused_args)
if properties.get('is_streaming', False):
# TODO(ccy): right now, if we pull this in when not in the worker
# environment, this will fail on not being able to pull in the correct gRPC
# C dependencies. I am investigating a fix.
from google.cloud.dataflow.worker import streamingworker # pylint: disable=g-import-not-at-top
# Initialize the random number generator, which is used to generate Windmill
# client IDs.
random.seed()
logging.info('Starting streaming worker.')
streamingworker.StreamingWorker(properties).run()
else:
logging.info('Starting batch worker.')
batchworker.BatchWorker(properties, sdk_pipeline_options).run()
def test_worker_fails_for_deferred_exceptions(self, mock_report_status,
mock_stop, mock_start,
mock_execute):
worker = batchworker.BatchWorker(self.dummy_properties(), {})
mock_work_item = mock.MagicMock()
worker.do_work(mock_work_item,
deferred_exception_details='deferred_exc')
mock_report_status.assert_called_with(completed=True,
exception_details='deferred_exc')
assert not mock_stop.called
assert not mock_start.called
assert not mock_execute.called
def test_worker_starts_and_stops_progress_reporter(self,
mock_report_status,
mock_stop, mock_start,
mock_execute):
worker = batchworker.BatchWorker(self.dummy_properties(), {})
mock_work_item = mock.MagicMock()
worker.do_work(mock_work_item)
mock_report_status.assert_called_with(completed=True,
exception_details=None)
mock_start.assert_called_once_with()
mock_execute.assert_called_once_with(mock.ANY)
mock_stop.assert_called_once_with()
def test_worker_requests_for_work(self, mock_lease_work, mock_workitem):
worker = batchworker.BatchWorker(self.dummy_properties(), {})
rand_work = object()
mock_lease_work.return_value = rand_work
mock_workitem.get_work_items.return_value = None
thread = threading.Thread(target=worker.run)
thread.start()
time.sleep(5)
worker.shutdown()
mock_lease_work.assert_called_with(
mock.ANY, worker.default_desired_lease_duration())
mock_workitem.get_work_items.assert_called_with(
rand_work, mock.ANY, mock.ANY)
def _run_send_completion_test(self, mock_report_status, mock_stop,
mock_start, mock_execute,
expected_exception):
worker = batchworker.BatchWorker(self.dummy_properties(), {})
mock_work_item = mock.MagicMock()
worker.do_work(mock_work_item)
class AnyStringWith(str):
def __eq__(self, other):
return self in other
mock_report_status.assert_called_with(
completed=True,
exception_details=AnyStringWith(expected_exception))
mock_start.assert_called_once_with()
mock_execute.assert_called_once_with(mock.ANY)
mock_stop.assert_called_once_with()