def test_notify_request_with_tq_batch_mode_false(self):
request = _gen_request()
result_summary = task_scheduler.schedule_request(request, None)
self.cfg.enable_batch_es_notifications = True
self.assertEqual(1, self.execute_tasks())
self._setup_client()
# Since use_tq is false, the requests below should be sent out immediately.
external_scheduler.notify_requests(self.cfg_foo,
[(request, result_summary)],
False,
False,
batch_mode=True)
external_scheduler.notify_requests(self.cfg_hoe,
[(request, result_summary)],
False,
False,
batch_mode=True)
called_with = self._client.called_with_requests
self.assertEqual(len(called_with), 2)
called_with.sort(key=lambda x: x.scheduler_id)
# Should have 1 notification and its id is foo.
self.assertEqual(len(called_with[0].notifications), 1)
self.assertEqual(called_with[0].scheduler_id, u'foo')
# Should have 1 notification and its id is hoe.
self.assertEqual(len(called_with[1].notifications), 1)
self.assertEqual(called_with[1].scheduler_id, u'hoe')
def test_notify_requests(self):
request = _gen_request()
result_summary = task_scheduler.schedule_request(request, None)
external_scheduler.notify_requests(self.es_cfg,
[(request, result_summary)], False,
False)
self.assertEqual(len(self._client.called_with_requests), 1)
called_with = self._client.called_with_requests[0]
self.assertEqual(len(called_with.notifications), 1)
notification = called_with.notifications[0]
self.assertEqual(request.created_ts,
notification.task.enqueued_time.ToDatetime())
self.assertEqual(request.task_id, notification.task.id)
self.assertEqual(request.num_task_slices,
len(notification.task.slices))
self.execute_tasks()
def test_notify_request_with_tq(self):
request = _gen_request()
result_summary = task_scheduler.schedule_request(request, None)
external_scheduler.notify_requests(self.es_cfg,
[(request, result_summary)], True,
False)
# There should have been no call to _get_client yet.
self.assertEqual(self._client, None)
self.execute_tasks()
# After taskqueue executes, there should be a call to the client.
self.assertEqual(len(self._client.called_with_requests), 1)
called_with = self._client.called_with_requests[0]
self.assertEqual(len(called_with.notifications), 1)
notification = called_with.notifications[0]
self.assertEqual(request.created_ts,
notification.task.enqueued_time.ToDatetime())
self.assertEqual(request.task_id, notification.task.id)
self.assertEqual(request.num_task_slices,
len(notification.task.slices))
def test_notify_request_with_tq_batch_mode(self):
request = _gen_request()
result_summary = task_scheduler.schedule_request(request, None)
self.assertEqual(1, self.execute_tasks())
# Create requests with different scheduler IDs.
external_scheduler.notify_requests(self.cfg_foo,
[(request, result_summary)],
True,
False,
batch_mode=True)
external_scheduler.notify_requests(self.cfg_foo,
[(request, result_summary)],
True,
False,
batch_mode=True)
external_scheduler.notify_requests(self.cfg_hoe,
[(request, result_summary)],
True,
False,
batch_mode=True)
self._setup_client()
# There should have been no call in _get_client yet.
self.assertEqual(len(self._client.called_with_requests), 0)
# Execute the kicker to call the pull queue worker.
# 3 tasks(kickers) were added to es-notify-kickers.
# 2 tasks(batched request) will be added to es-notify-tasks
# once the kicker is done.
self.assertEqual(5, self.execute_tasks())
called_with = self._client.called_with_requests
# There should have 2 calls to the external scheduler.
self.assertEqual(len(called_with), 2)
called_with.sort(key=lambda x: x.scheduler_id)
# Request foo should have 2 notifications.
self.assertEqual(len(called_with[0].notifications), 2)
self.assertEqual(called_with[0].scheduler_id, u'foo')
# Request hoe should have 1 notification.
self.assertEqual(len(called_with[1].notifications), 1)
self.assertEqual(called_with[1].scheduler_id, u'hoe')
# There should be no task remained in the pull queue.
stats = taskqueue.QueueStatistics.fetch('es-notify-tasks-batch')
self.assertEqual(0, stats.tasks)