]>
Commit | Line | Data |
---|---|---|
11fdf7f2 TL |
1 | # -*- coding: utf-8 -*- |
2 | # pylint: disable=too-many-arguments | |
3 | from __future__ import absolute_import | |
4 | ||
5 | import json | |
6 | import threading | |
7 | import time | |
8 | ||
9 | import cherrypy | |
10 | from cherrypy._cptools import HandlerWrapperTool | |
11 | from cherrypy.test import helper | |
12 | ||
13 | from mgr_module import CLICommand, MgrModule | |
14 | ||
15 | from .. import logger, mgr | |
16 | from ..controllers import json_error_page, generate_controller_routes | |
17 | from ..services.auth import AuthManagerTool | |
18 | from ..services.exception import dashboard_exception_handler | |
19 | ||
20 | ||
21 | class CmdException(Exception): | |
22 | def __init__(self, retcode, message): | |
23 | super(CmdException, self).__init__(message) | |
24 | self.retcode = retcode | |
25 | ||
26 | ||
27 | def exec_dashboard_cmd(command_handler, cmd, **kwargs): | |
28 | cmd_dict = {'prefix': 'dashboard {}'.format(cmd)} | |
29 | cmd_dict.update(kwargs) | |
30 | if cmd_dict['prefix'] not in CLICommand.COMMANDS: | |
31 | ret, out, err = command_handler(cmd_dict) | |
32 | if ret < 0: | |
33 | raise CmdException(ret, err) | |
34 | try: | |
35 | return json.loads(out) | |
36 | except ValueError: | |
37 | return out | |
38 | ||
39 | ret, out, err = CLICommand.COMMANDS[cmd_dict['prefix']].call(mgr, cmd_dict, | |
40 | None) | |
41 | if ret < 0: | |
42 | raise CmdException(ret, err) | |
43 | try: | |
44 | return json.loads(out) | |
45 | except ValueError: | |
46 | return out | |
47 | ||
48 | ||
49 | class KVStoreMockMixin(object): | |
50 | CONFIG_KEY_DICT = {} | |
51 | ||
52 | @classmethod | |
53 | def mock_set_module_option(cls, attr, val): | |
54 | cls.CONFIG_KEY_DICT[attr] = val | |
55 | ||
56 | @classmethod | |
57 | def mock_get_module_option(cls, attr, default=None): | |
58 | return cls.CONFIG_KEY_DICT.get(attr, default) | |
59 | ||
60 | @classmethod | |
61 | def mock_kv_store(cls): | |
62 | cls.CONFIG_KEY_DICT.clear() | |
63 | mgr.set_module_option.side_effect = cls.mock_set_module_option | |
64 | mgr.get_module_option.side_effect = cls.mock_get_module_option | |
65 | # kludge below | |
66 | mgr.set_store.side_effect = cls.mock_set_module_option | |
67 | mgr.get_store.side_effect = cls.mock_get_module_option | |
68 | ||
69 | @classmethod | |
70 | def get_key(cls, key): | |
71 | return cls.CONFIG_KEY_DICT.get(key, None) | |
72 | ||
73 | ||
74 | class CLICommandTestMixin(KVStoreMockMixin): | |
75 | @classmethod | |
76 | def exec_cmd(cls, cmd, **kwargs): | |
77 | return exec_dashboard_cmd(None, cmd, **kwargs) | |
78 | ||
79 | ||
80 | class ControllerTestCase(helper.CPWebCase): | |
81 | @classmethod | |
82 | def setup_controllers(cls, ctrl_classes, base_url=''): | |
83 | if not isinstance(ctrl_classes, list): | |
84 | ctrl_classes = [ctrl_classes] | |
85 | mapper = cherrypy.dispatch.RoutesDispatcher() | |
86 | endpoint_list = [] | |
87 | for ctrl in ctrl_classes: | |
88 | inst = ctrl() | |
89 | for endpoint in ctrl.endpoints(): | |
90 | endpoint.inst = inst | |
91 | endpoint_list.append(endpoint) | |
92 | endpoint_list = sorted(endpoint_list, key=lambda e: e.url) | |
93 | for endpoint in endpoint_list: | |
94 | generate_controller_routes(endpoint, mapper, base_url) | |
95 | if base_url == '': | |
96 | base_url = '/' | |
97 | cherrypy.tree.mount(None, config={ | |
98 | base_url: {'request.dispatch': mapper}}) | |
99 | ||
100 | def __init__(self, *args, **kwargs): | |
101 | cherrypy.tools.authenticate = AuthManagerTool() | |
102 | cherrypy.tools.dashboard_exception_handler = HandlerWrapperTool(dashboard_exception_handler, | |
103 | priority=31) | |
104 | cherrypy.config.update({ | |
105 | 'error_page.default': json_error_page, | |
106 | 'tools.json_in.on': True, | |
107 | 'tools.json_in.force': False | |
108 | }) | |
109 | super(ControllerTestCase, self).__init__(*args, **kwargs) | |
110 | ||
111 | def _request(self, url, method, data=None): | |
112 | if not data: | |
113 | b = None | |
114 | h = None | |
115 | else: | |
116 | b = json.dumps(data) | |
117 | h = [('Content-Type', 'application/json'), | |
118 | ('Content-Length', str(len(b)))] | |
119 | self.getPage(url, method=method, body=b, headers=h) | |
120 | ||
121 | def _get(self, url): | |
122 | self._request(url, 'GET') | |
123 | ||
124 | def _post(self, url, data=None): | |
125 | self._request(url, 'POST', data) | |
126 | ||
127 | def _delete(self, url, data=None): | |
128 | self._request(url, 'DELETE', data) | |
129 | ||
130 | def _put(self, url, data=None): | |
131 | self._request(url, 'PUT', data) | |
132 | ||
133 | def _task_request(self, method, url, data, timeout): | |
134 | self._request(url, method, data) | |
135 | if self.status != '202 Accepted': | |
136 | logger.info("task finished immediately") | |
137 | return | |
138 | ||
139 | res = self.jsonBody() | |
140 | self.assertIsInstance(res, dict) | |
141 | self.assertIn('name', res) | |
142 | self.assertIn('metadata', res) | |
143 | ||
144 | task_name = res['name'] | |
145 | task_metadata = res['metadata'] | |
146 | ||
147 | # pylint: disable=protected-access | |
148 | class Waiter(threading.Thread): | |
149 | def __init__(self, task_name, task_metadata, tc): | |
150 | super(Waiter, self).__init__() | |
151 | self.task_name = task_name | |
152 | self.task_metadata = task_metadata | |
153 | self.ev = threading.Event() | |
154 | self.abort = False | |
155 | self.res_task = None | |
156 | self.tc = tc | |
157 | ||
158 | def run(self): | |
159 | running = True | |
160 | while running and not self.abort: | |
161 | logger.info("task (%s, %s) is still executing", self.task_name, | |
162 | self.task_metadata) | |
163 | time.sleep(1) | |
164 | self.tc._get('/api/task?name={}'.format(self.task_name)) | |
165 | res = self.tc.jsonBody() | |
166 | for task in res['finished_tasks']: | |
167 | if task['metadata'] == self.task_metadata: | |
168 | # task finished | |
169 | running = False | |
170 | self.res_task = task | |
171 | self.ev.set() | |
172 | ||
173 | thread = Waiter(task_name, task_metadata, self) | |
174 | thread.start() | |
175 | status = thread.ev.wait(timeout) | |
176 | if not status: | |
177 | # timeout expired | |
178 | thread.abort = True | |
179 | thread.join() | |
180 | raise Exception("Waiting for task ({}, {}) to finish timed out" | |
181 | .format(task_name, task_metadata)) | |
182 | logger.info("task (%s, %s) finished", task_name, task_metadata) | |
183 | if thread.res_task['success']: | |
184 | self.body = json.dumps(thread.res_task['ret_value']) | |
185 | if method == 'POST': | |
186 | self.status = '201 Created' | |
187 | elif method == 'PUT': | |
188 | self.status = '200 OK' | |
189 | elif method == 'DELETE': | |
190 | self.status = '204 No Content' | |
191 | return | |
192 | else: | |
193 | if 'status' in thread.res_task['exception']: | |
194 | self.status = thread.res_task['exception']['status'] | |
195 | else: | |
196 | self.status = 500 | |
197 | self.body = json.dumps(thread.res_task['exception']) | |
198 | return | |
199 | ||
200 | def _task_post(self, url, data=None, timeout=60): | |
201 | self._task_request('POST', url, data, timeout) | |
202 | ||
203 | def _task_delete(self, url, timeout=60): | |
204 | self._task_request('DELETE', url, None, timeout) | |
205 | ||
206 | def _task_put(self, url, data=None, timeout=60): | |
207 | self._task_request('PUT', url, data, timeout) | |
208 | ||
209 | def jsonBody(self): | |
210 | body_str = self.body.decode('utf-8') if isinstance(self.body, bytes) else self.body | |
211 | return json.loads(body_str) | |
212 | ||
213 | def assertJsonBody(self, data, msg=None): | |
214 | """Fail if value != self.body.""" | |
215 | json_body = self.jsonBody() | |
216 | if data != json_body: | |
217 | if msg is None: | |
218 | msg = 'expected body:\n%r\n\nactual body:\n%r' % ( | |
219 | data, json_body) | |
220 | self._handlewebError(msg) | |
221 | ||
222 | def assertInJsonBody(self, data, msg=None): | |
223 | json_body = self.jsonBody() | |
224 | if data not in json_body: | |
225 | if msg is None: | |
226 | msg = 'expected %r to be in %r' % (data, json_body) | |
227 | self._handlewebError(msg) |