blob: 3c7305ff7b1c32f5b8add61b18d36a67683fefa5 [file] [log] [blame]
Xixuan Wu835dee22017-09-07 10:47:29 -07001# Copyright 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Module for task_executor unittests."""
6# pylint: disable=g-bad-import-order
Xinan Lin3ba18a02019-08-13 15:44:55 -07007# pylint: disable=unused-argument
Xixuan Wu835dee22017-09-07 10:47:29 -07008
Xinan Lin9e4917d2019-11-04 10:58:47 -08009import collections
Xixuan Wu835dee22017-09-07 10:47:29 -070010import mock
11import os
12import sys
13import unittest
14
15import task_executor
16
17from google.appengine.api import taskqueue
18from google.appengine.ext import testbed
19
20
Xinan Lin3ba18a02019-08-13 15:44:55 -070021class FakeFrontdoorClient(object):
22
23 # pylint: disable=g-deprecated-member-used
24 def __init__(self, success_num=sys.maxint, error_num=0):
25 self.success_num = success_num
26 self.error_num = error_num
Xinan Lin9e4917d2019-11-04 10:58:47 -080027 self.frontdoor_run_count = collections.defaultdict(list)
Xinan Lin3ba18a02019-08-13 15:44:55 -070028
Xinan Lin9e4917d2019-11-04 10:58:47 -080029 def multirequest_run(self, tasks, suite):
30 executed_tasks = []
31 for task in tasks:
32 params = task.extract_params()
33 num = int(params.get('num', 0))
34 if num > self.success_num and num <= self.success_num + self.error_num:
35 raise ValueError('test')
36 executed_tasks.append(task)
37 self.frontdoor_run_count[suite].append(executed_tasks)
38 return executed_tasks
Xinan Lin3ba18a02019-08-13 15:44:55 -070039
40
Xixuan Wu835dee22017-09-07 10:47:29 -070041class TaskExecutorTestCase(unittest.TestCase):
42
43 def setUp(self):
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080044 super(TaskExecutorTestCase, self).setUp()
Xixuan Wu835dee22017-09-07 10:47:29 -070045 self.testbed = testbed.Testbed()
46 self.testbed.activate()
47 self.addCleanup(self.testbed.deactivate)
48
49 # root_path must be set the location of queue.yaml.
50 # Otherwise, only the 'default' queue will be available.
51 self.testbed.init_taskqueue_stub(
52 root_path=os.path.join(os.path.dirname(__file__)))
53 self.taskqueue_stub = self.testbed.get_stub(
54 testbed.TASKQUEUE_SERVICE_NAME)
55
Prathmesh Prabhu40f1e1e2020-03-07 00:04:41 -080056 p = mock.patch('global_config.GAE_TESTING', return_value=True)
57 p.start()
58 self.addCleanup(p.stop)
59
Xixuan Wu835dee22017-09-07 10:47:29 -070060 def testPushTask(self):
61 suite_kwargs = {'suite': 'fake_suite'}
62 task_executor.push(task_executor.SUITES_QUEUE, **suite_kwargs)
63 tasks = self.taskqueue_stub.get_filtered_tasks()
64 self.assertEqual(len(tasks), 1)
65 self.assertEqual(suite_kwargs, tasks[0].extract_params())
66
Xinan Lin3ba18a02019-08-13 15:44:55 -070067 def testBatchExecuteTask(self):
Xinan Lina18d9882019-12-18 11:16:33 -080068 """Test task_executor executes tasks not greater than the BATCH_SIZE."""
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080069 batch_size = 100
Xinan Lin9e4917d2019-11-04 10:58:47 -080070 extra_num = 10
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080071 for i in range(batch_size / 2):
Xinan Lin9e4917d2019-11-04 10:58:47 -080072 suite_kwargs = {'suite': 'foo_suite',
Xinan Lin3ba18a02019-08-13 15:44:55 -070073 'num': i + 1}
Xinan Lin9e4917d2019-11-04 10:58:47 -080074 task_executor.push(task_executor.SUITES_QUEUE,
75 tag='foo_suite', **suite_kwargs)
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080076 for i in range(batch_size / 2 + extra_num):
Xinan Lin9e4917d2019-11-04 10:58:47 -080077 suite_kwargs = {'suite': 'hoo_suite',
78 'num': i + 1}
79 task_executor.push(task_executor.SUITES_QUEUE,
80 tag='hoo_suite', **suite_kwargs)
Xinan Lin3ba18a02019-08-13 15:44:55 -070081 with mock.patch('buildbucket.TestPlatformClient',
82 return_value=FakeFrontdoorClient()):
Xinan Lin9e4917d2019-11-04 10:58:47 -080083 with mock.patch('global_config.GAE_TESTING', return_value=True):
84 task_processor = task_executor.TaskProcessor(
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080085 task_executor.SUITES_QUEUE,
86 options=task_executor.Options(
87 batch_size=batch_size, multirequest_size=1))
Xinan Lin9e4917d2019-11-04 10:58:47 -080088 task_processor.batch_execute()
Xinan Lina18d9882019-12-18 11:16:33 -080089 # After batch_execute runs, the extra tasks should remain in the queue.
Xinan Lin9e4917d2019-11-04 10:58:47 -080090 tasks = self.taskqueue_stub.get_filtered_tasks()
91 self.assertEqual(extra_num, len(tasks))
Xixuan Wua5a29442017-10-11 11:03:02 -070092
Xinan Lin9e4917d2019-11-04 10:58:47 -080093 def testBatchExecuteTaskFailedFrontdoorTotally(self):
Xixuan Wu835dee22017-09-07 10:47:29 -070094 """Test task_executor fails at the beginning, and no tasks are deleted."""
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080095 batch_size = 100
96 multirequest_size = 30
Xixuan Wu835dee22017-09-07 10:47:29 -070097 extra_num = 10
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -080098 suite_kwargs = {'suite': 'fake_suite'}
99 for i in range(multirequest_size + extra_num):
Xixuan Wu835dee22017-09-07 10:47:29 -0700100 suite_kwargs['num'] = i + 1
101 task_executor.push(task_executor.SUITES_QUEUE, **suite_kwargs)
102
103 # Before batch_execute
104 tasks = self.taskqueue_stub.get_filtered_tasks()
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800105 self.assertEqual(len(tasks), multirequest_size + extra_num)
Xinan Lin9e4917d2019-11-04 10:58:47 -0800106 with mock.patch('buildbucket.TestPlatformClient',
107 return_value=FakeFrontdoorClient(0, error_num=len(tasks))):
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800108 task_processor = task_executor.TaskProcessor(
109 task_executor.SUITES_QUEUE,
110 options=task_executor.Options(
111 batch_size=batch_size, multirequest_size=multirequest_size))
Xixuan Wu0a8d3ee2017-10-19 11:33:26 -0700112 task_processor.batch_execute()
113 # After batch_execute, no tasks are deleted from task queue, due
114 # to they're all failed to kick off.
Xixuan Wu835dee22017-09-07 10:47:29 -0700115 tasks = self.taskqueue_stub.get_filtered_tasks()
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800116 self.assertEqual(len(tasks), multirequest_size + extra_num)
Xixuan Wu835dee22017-09-07 10:47:29 -0700117
Xinan Lin9e4917d2019-11-04 10:58:47 -0800118 def testBatchExecuteTaskFailedFrontdoorPartially(self):
Xixuan Wu835dee22017-09-07 10:47:29 -0700119 """Test task_executor fails halfway, and only executed tasks are deleted."""
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800120 # Batch large enough to cover all added tasks.
121 batch_size = 1000
122 multirequest_size = 30
Xixuan Wu835dee22017-09-07 10:47:29 -0700123 extra_num = 10
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800124 suite_kwargs = {'suite': 'fake_suite'}
125 success_num = (multirequest_size + extra_num) / 2
126 error_num = multirequest_size + extra_num
127 for i in range(multirequest_size + extra_num):
Xixuan Wu835dee22017-09-07 10:47:29 -0700128 suite_kwargs['num'] = i + 1
129 task_executor.push(task_executor.SUITES_QUEUE, **suite_kwargs)
130
131 # Before batch_execute
132 tasks = self.taskqueue_stub.get_filtered_tasks()
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800133 self.assertEqual(len(tasks), multirequest_size + extra_num)
Xinan Lin9e4917d2019-11-04 10:58:47 -0800134 with mock.patch('buildbucket.TestPlatformClient',
135 return_value=FakeFrontdoorClient(success_num, error_num)):
Prathmesh Prabhu7b961d52020-03-06 23:57:42 -0800136 task_processor = task_executor.TaskProcessor(
137 task_executor.SUITES_QUEUE,
138 options=task_executor.Options(
139 batch_size=batch_size, multirequest_size=multirequest_size))
Xixuan Wu0a8d3ee2017-10-19 11:33:26 -0700140 task_processor.batch_execute()
141 # After batch_execute, only failed suites and extra suites are
142 # kept in task queue.
Xixuan Wu835dee22017-09-07 10:47:29 -0700143 tasks = self.taskqueue_stub.get_filtered_tasks()
Xinan Lin9e4917d2019-11-04 10:58:47 -0800144 self.assertEqual(len(tasks), error_num)
Xixuan Wu835dee22017-09-07 10:47:29 -0700145
146 def testBatchExecuteTaskFailedLeasing(self):
147 """Test task_executor fails to lease task."""
148 suite_kwargs = {'suite': 'fake_suite'}
149 task_executor.push(task_executor.SUITES_QUEUE, **suite_kwargs)
150
Xinan Lin9e4917d2019-11-04 10:58:47 -0800151 with mock.patch('buildbucket.TestPlatformClient',
152 return_value=FakeFrontdoorClient(False)):
153 task_processor = task_executor.TaskProcessor('nonExistentQueue')
Xixuan Wu835dee22017-09-07 10:47:29 -0700154 self.assertRaises(taskqueue.UnknownQueueError,
155 task_processor.batch_execute)
156 # After batch_execute fails, no tasks are deleted from task queue.
157 tasks = self.taskqueue_stub.get_filtered_tasks()
158 self.assertEqual(len(tasks), 1)
159
160
161if __name__ == '__main__':
162 unittest.main()