test_worker.py 17 KB
Newer Older
Philip ABBET's avatar
Philip ABBET committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :

###############################################################################
#                                                                             #
# Copyright (c) 2017 Idiap Research Institute, http://www.idiap.ch/           #
# Contact: beat.support@idiap.ch                                              #
#                                                                             #
# This file is part of the beat.core module of the BEAT platform.             #
#                                                                             #
# Commercial License Usage                                                    #
# Licensees holding valid commercial BEAT licenses may use this file in       #
# accordance with the terms contained in a written agreement between you      #
# and Idiap. For further information contact tto@idiap.ch                     #
#                                                                             #
# Alternatively, this file may be used under the terms of the GNU Affero      #
# Public License version 3 as published by the Free Software and appearing    #
# in the file LICENSE.AGPL included in the packaging of this file.            #
# The BEAT platform is distributed in the hope that it will be useful, but    #
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY  #
# or FITNESS FOR A PARTICULAR PURPOSE.                                        #
#                                                                             #
# You should have received a copy of the GNU Affero Public License along      #
# with the BEAT platform. If not, see http://www.gnu.org/licenses/.           #
#                                                                             #
###############################################################################


# Tests for experiment execution

import os
Philip ABBET's avatar
Philip ABBET committed
32

Philip ABBET's avatar
Philip ABBET committed
33
34
35
36
37
import logging
logger = logging.getLogger(__name__)

import unittest
import simplejson
Philip ABBET's avatar
Philip ABBET committed
38
import multiprocessing
39
import Queue
Philip ABBET's avatar
Philip ABBET committed
40
from time import time
41
from time import sleep
Philip ABBET's avatar
Philip ABBET committed
42
43
44
45
46
47
48
49

from ..scripts import worker
from ..worker import WorkerController
from ..dock import Host

from . import prefix, tmp_prefix


Philip ABBET's avatar
Philip ABBET committed
50
51
52
#----------------------------------------------------------


Philip ABBET's avatar
Philip ABBET committed
53
54
55
56
WORKER1 = 'worker1'
WORKER2 = 'worker2'


Philip ABBET's avatar
Philip ABBET committed
57
58
59
#----------------------------------------------------------


Philip ABBET's avatar
Philip ABBET committed
60
CONFIGURATION1 = {
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    'queue': 'queue',
    'inputs': {
        'in_data': {
            'set': 'double',
            'protocol': 'double',
            'database': 'integers_db/1',
            'output': 'a',
            'path': 'ec/89/e5/6e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55',
            'endpoint': 'a',
            'hash': 'ec89e56e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55',
            'channel': 'integers'
        }
    },
    'algorithm': 'user/integers_echo/1',
    'parameters': {},
    'environment': {
        'name': 'Python 2.7',
        'version': '1.2.0'
    },
    'outputs': {
        'out_data': {
            'path': '20/61/b6/2df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681',
            'endpoint': 'out_data',
            'hash': '2061b62df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681',
            'channel': 'integers'
        }
    },
    'nb_slots': 1,
    'channel': 'integers'
Philip ABBET's avatar
Philip ABBET committed
90
91
92
}


Philip ABBET's avatar
Philip ABBET committed
93
94
95
#----------------------------------------------------------


Philip ABBET's avatar
Philip ABBET committed
96
CONFIGURATION2 = {
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    'queue': 'queue',
    'inputs': {
        'in_data': {
            'set': 'double',
            'protocol': 'double',
            'database': 'integers_db/1',
            'output': 'a',
            'path': 'ec/89/e5/6e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55',
            'endpoint': 'a',
            'hash': 'ec89e56e161d2cb012ef6ac8acf59bf453a6328766f90dc9baba9eb14ea23c55',
            'channel': 'integers'
        }
    },
    'algorithm': 'user/integers_echo/1',
    'parameters': {},
    'environment': {
        'name': 'Python 2.7',
        'version': '1.2.0'
    },
    'outputs': {
        'out_data': {
            'path': '40/61/b6/2df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681',
            'endpoint': 'out_data',
            'hash': '4061b62df3c3bedd5366f4a625c5d87ffbf5a26007c46c456e9abf21b46c6681',
            'channel': 'integers'
        }
    },
    'nb_slots': 1,
    'channel': 'integers'
Philip ABBET's avatar
Philip ABBET committed
126
127
128
}


Philip ABBET's avatar
Philip ABBET committed
129
130
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
131

132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
class ControllerProcess(multiprocessing.Process):

    def __init__(self, queue):
        super(ControllerProcess, self).__init__()

        self.queue = queue


    def run(self):
        self.queue.put('STARTED')

        def onWorkerReady(name):
            self.queue.put('READY ' + name)

        def onWorkerGone(name):
            self.queue.put('GONE ' + name)

        self.controller = WorkerController(
            '127.0.0.1',
            port=51000,
            callbacks=dict(
                onWorkerReady = onWorkerReady,
                onWorkerGone = onWorkerGone,
            )
        )

        while True:
            self.controller.process(100)

            try:
                command = self.queue.get_nowait()
                if command == 'STOP':
                    break
            except Queue.Empty:
                pass

        self.controller.destroy()


#----------------------------------------------------------


Philip ABBET's avatar
Philip ABBET committed
174
class WorkerProcess(multiprocessing.Process):
Philip ABBET's avatar
Philip ABBET committed
175

Philip ABBET's avatar
Philip ABBET committed
176
177
    def __init__(self, queue, arguments):
        super(WorkerProcess, self).__init__()
Philip ABBET's avatar
Philip ABBET committed
178

Philip ABBET's avatar
Philip ABBET committed
179
180
        self.queue = queue
        self.arguments = arguments
Philip ABBET's avatar
Philip ABBET committed
181
182


Philip ABBET's avatar
Philip ABBET committed
183
184
185
    def run(self):
        self.queue.put('STARTED')
        worker.main(self.arguments)
Philip ABBET's avatar
Philip ABBET committed
186
187


Philip ABBET's avatar
Philip ABBET committed
188
189
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
190

191
class TestWorkerBase(unittest.TestCase):
Philip ABBET's avatar
Philip ABBET committed
192

Philip ABBET's avatar
Philip ABBET committed
193
    def __init__(self, methodName='runTest'):
194
        super(TestWorkerBase, self).__init__(methodName)
Philip ABBET's avatar
Philip ABBET committed
195
        self.controller = None
196
197
        self.connected_workers = []
        self.worker_processes = {}
Philip ABBET's avatar
Philip ABBET committed
198
        self.docker = False
Philip ABBET's avatar
Philip ABBET committed
199
200


Philip ABBET's avatar
Philip ABBET committed
201
    def setUp(self):
202
203
204
205
206
207
208
209
210
211
212
213
214
        self.shutdown_everything()  # In case another test failed badly during its setUp()


    def tearDown(self):
        self.shutdown_everything()


    def shutdown_everything(self):
        for name in list(self.worker_processes.keys()):
            self.stop_worker(name)

        self.worker_processes = {}
        self.connected_workers = []
Philip ABBET's avatar
Philip ABBET committed
215

216
217
218
219
220
        self.stop_controller()


    def start_controller(self, port=None):
        self.connected_workers = []
Philip ABBET's avatar
Philip ABBET committed
221

Philip ABBET's avatar
Philip ABBET committed
222
        def onWorkerReady(name):
223
224
225
226
            self.connected_workers.append(name)

        def onWorkerGone(name):
            self.connected_workers.remove(name)
Philip ABBET's avatar
Philip ABBET committed
227

Philip ABBET's avatar
Philip ABBET committed
228
        self.controller = WorkerController(
229
230
231
232
233
234
            '127.0.0.1',
            port=port,
            callbacks=dict(
                onWorkerReady = onWorkerReady,
                onWorkerGone = onWorkerGone,
            )
Philip ABBET's avatar
Philip ABBET committed
235
        )
Philip ABBET's avatar
Philip ABBET committed
236

237
238
239
240
241
242
243
244
245
246
        self.controller.process(100)


    def stop_controller(self):
        if self.controller is not None:
            self.controller.destroy()
            self.controller = None


    def start_worker(self, name, address=None):
Philip ABBET's avatar
Philip ABBET committed
247
248
249
        args = [
          '--prefix=%s' % prefix,
          '--cache=%s' % tmp_prefix,
250
          '--name=%s' % name,
251
          # '-vv',
252
          self.controller.address if address is None else address,
Philip ABBET's avatar
Philip ABBET committed
253
        ]
Philip ABBET's avatar
Philip ABBET committed
254

Philip ABBET's avatar
Philip ABBET committed
255
256
        if self.docker:
            args.insert(3, '--docker')
Philip ABBET's avatar
Philip ABBET committed
257

258
259
260
261
262
263
        worker_process = WorkerProcess(multiprocessing.Queue(), args)
        worker_process.start()

        worker_process.queue.get()

        self.worker_processes[name] = worker_process
Philip ABBET's avatar
Philip ABBET committed
264
265


266
267
268
269
270
271
272
273
    def stop_worker(self, name):
        if name in self.worker_processes:
            self.worker_processes[name].terminate()
            self.worker_processes[name].join()
            del self.worker_processes[name]


    def wait_for_worker_connection(self, name):
Philip ABBET's avatar
Philip ABBET committed
274
        start = time()
275
        while name not in self.connected_workers:
Philip ABBET's avatar
Philip ABBET committed
276
277
            self.assertTrue(self.controller.process(100) is None)
            self.assertTrue(time() - start < 10)  # Exit after 10 seconds
Philip ABBET's avatar
Philip ABBET committed
278

279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        self.assertTrue(name in self.controller.workers)


    def wait_for_worker_disconnection(self, name):
        start = time()
        while name in self.connected_workers:
            self.assertTrue(self.controller.process(100) is None)
            self.assertTrue(time() - start < 10)  # Exit after 10 seconds

        self.assertTrue(name not in self.controller.workers)


#----------------------------------------------------------


class TestConnection(TestWorkerBase):

    def test_worker_connection(self):
        self.start_controller()

        self.assertEqual(len(self.connected_workers), 0)
        self.assertEqual(len(self.controller.workers), 0)

        self.start_worker(WORKER1)

        self.wait_for_worker_connection(WORKER1)

        self.assertEqual(len(self.connected_workers), 1)
Philip ABBET's avatar
Philip ABBET committed
307
        self.assertEqual(len(self.controller.workers), 1)
Philip ABBET's avatar
Philip ABBET committed
308

Philip ABBET's avatar
Philip ABBET committed
309

310
311
312
    def test_worker_disconnection(self):
        self.start_controller()
        self.start_worker(WORKER1)
Philip ABBET's avatar
Philip ABBET committed
313

314
        self.wait_for_worker_connection(WORKER1)
Philip ABBET's avatar
Philip ABBET committed
315

Philip ABBET's avatar
Philip ABBET committed
316
317
        sleep(1)

318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
        self.stop_worker(WORKER1)

        self.wait_for_worker_disconnection(WORKER1)


    def test_two_workers_connection(self):
        self.start_controller()

        self.assertEqual(len(self.connected_workers), 0)
        self.assertEqual(len(self.controller.workers), 0)

        self.start_worker(WORKER1)
        self.start_worker(WORKER2)

        self.wait_for_worker_connection(WORKER1)
        self.wait_for_worker_connection(WORKER2)

        self.assertEqual(len(self.connected_workers), 2)
        self.assertEqual(len(self.controller.workers), 2)


    def test_scheduler_last(self):
        self.start_worker(WORKER1, address='tcp://127.0.0.1:51000')
        sleep(1)

        self.start_controller(port=51000)

        self.wait_for_worker_connection(WORKER1)


348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    def test_scheduler_shutdown(self):
        controller = ControllerProcess(multiprocessing.Queue())
        controller.start()

        message = controller.queue.get()
        self.assertEqual(message, 'STARTED')

        self.start_worker(WORKER1, 'tcp://127.0.0.1:51000')

        message = controller.queue.get()
        self.assertEqual(message, 'READY ' + WORKER1)

        controller.queue.put('STOP')

        sleep(1)

        controller = ControllerProcess(multiprocessing.Queue())
        controller.start()

        message = controller.queue.get()
        self.assertEqual(message, 'STARTED')

        message = controller.queue.get()
        self.assertEqual(message, 'READY ' + WORKER1)

        controller.queue.put('STOP')


376
377
378
379
380
381
382
383
384
385
386
387
#----------------------------------------------------------


class TestOneWorker(TestWorkerBase):


    def setUp(self):
        super(TestOneWorker, self).setUp()

        self.start_controller()
        self.start_worker(WORKER1)
        self.wait_for_worker_connection(WORKER1)
Philip ABBET's avatar
Philip ABBET committed
388
389


Philip ABBET's avatar
Philip ABBET committed
390
391
392
    def _wait(self, max=100):
        message = None
        nb = 0
Philip ABBET's avatar
Philip ABBET committed
393

Philip ABBET's avatar
Philip ABBET committed
394
395
396
        while (message is None) and (nb < max):
            message = self.controller.process(100)
            nb += 1
Philip ABBET's avatar
Philip ABBET committed
397

Philip ABBET's avatar
Philip ABBET committed
398
        return message
Philip ABBET's avatar
Philip ABBET committed
399
400


Philip ABBET's avatar
Philip ABBET committed
401
402
    def _check_done(self, message, expected_worker, expected_job_id):
        self.assertTrue(message is not None)
Philip ABBET's avatar
Philip ABBET committed
403

Philip ABBET's avatar
Philip ABBET committed
404
        (worker, status, job_id, data) = message
Philip ABBET's avatar
Philip ABBET committed
405

Philip ABBET's avatar
Philip ABBET committed
406
407
408
        self.assertEqual(worker, expected_worker)
        self.assertEqual(status, WorkerController.DONE)
        self.assertEqual(job_id, expected_job_id)
Philip ABBET's avatar
Philip ABBET committed
409

Philip ABBET's avatar
Philip ABBET committed
410
        result = simplejson.loads(data[0])
Philip ABBET's avatar
Philip ABBET committed
411

Philip ABBET's avatar
Philip ABBET committed
412
        self.assertEqual(result['status'], 0)
Philip ABBET's avatar
Philip ABBET committed
413
414


Philip ABBET's avatar
Philip ABBET committed
415
416
    def test_success(self):
        self.controller.execute(WORKER1, 1, CONFIGURATION1)
Philip ABBET's avatar
Philip ABBET committed
417

Philip ABBET's avatar
Philip ABBET committed
418
        message = self._wait()
Philip ABBET's avatar
Philip ABBET committed
419

Philip ABBET's avatar
Philip ABBET committed
420
        self._check_done(message, WORKER1, 1)
Philip ABBET's avatar
Philip ABBET committed
421
422


Philip ABBET's avatar
Philip ABBET committed
423
424
425
    def test_processing_error(self):
        config = dict(CONFIGURATION1)
        config['algorithm'] = 'user/integers_crash/1'
Philip ABBET's avatar
Philip ABBET committed
426

Philip ABBET's avatar
Philip ABBET committed
427
        self.controller.execute(WORKER1, 1, config)
Philip ABBET's avatar
Philip ABBET committed
428

Philip ABBET's avatar
Philip ABBET committed
429
        (worker, status, job_id, data) = self._wait()
Philip ABBET's avatar
Philip ABBET committed
430

Philip ABBET's avatar
Philip ABBET committed
431
        self.assertEqual(worker, WORKER1)
Philip ABBET's avatar
Philip ABBET committed
432
        self.assertEqual(status, WorkerController.JOB_ERROR)
Philip ABBET's avatar
Philip ABBET committed
433
        self.assertEqual(job_id, 1)
Philip ABBET's avatar
Philip ABBET committed
434

Philip ABBET's avatar
Philip ABBET committed
435
        result = simplejson.loads(data[0])
Philip ABBET's avatar
Philip ABBET committed
436

Philip ABBET's avatar
Philip ABBET committed
437
438
        self.assertEqual(result['status'], 1)
        self.assertTrue('a = b' in result['user_error'])
Philip ABBET's avatar
Philip ABBET committed
439
440


Philip ABBET's avatar
Philip ABBET committed
441
442
443
    def test_error_unknown_algorithm(self):
        config = dict(CONFIGURATION1)
        config['algorithm'] = 'user/unknown/1'
Philip ABBET's avatar
Philip ABBET committed
444

Philip ABBET's avatar
Philip ABBET committed
445
        self.controller.execute(WORKER1, 1, config)
Philip ABBET's avatar
Philip ABBET committed
446

Philip ABBET's avatar
Philip ABBET committed
447
        (worker, status, job_id, data) = self._wait()
Philip ABBET's avatar
Philip ABBET committed
448

Philip ABBET's avatar
Philip ABBET committed
449
450
451
452
        self.assertEqual(worker, WORKER1)
        self.assertEqual(status, WorkerController.JOB_ERROR)
        self.assertEqual(job_id, 1)
        self.assertTrue(len(data) > 0)
Philip ABBET's avatar
Philip ABBET committed
453
454


Philip ABBET's avatar
Philip ABBET committed
455
456
457
    def test_error_syntax_error(self):
        config = dict(CONFIGURATION1)
        config['algorithm'] = 'user/syntax_error/1'
Philip ABBET's avatar
Philip ABBET committed
458

Philip ABBET's avatar
Philip ABBET committed
459
        self.controller.execute(WORKER1, 1, config)
Philip ABBET's avatar
Philip ABBET committed
460

Philip ABBET's avatar
Philip ABBET committed
461
        (worker, status, job_id, data) = self._wait()
Philip ABBET's avatar
Philip ABBET committed
462

Philip ABBET's avatar
Philip ABBET committed
463
464
465
466
        self.assertEqual(worker, WORKER1)
        self.assertEqual(status, WorkerController.JOB_ERROR)
        self.assertEqual(job_id, 1)
        self.assertTrue(len(data) > 0)
Philip ABBET's avatar
Philip ABBET committed
467
468


Philip ABBET's avatar
Philip ABBET committed
469
    def test_multiple_jobs(self):
Philip ABBET's avatar
Philip ABBET committed
470
471
        config = dict(CONFIGURATION1)
        config['algorithm'] = 'user/integers_echo_slow/1'
Philip ABBET's avatar
Philip ABBET committed
472

Philip ABBET's avatar
Philip ABBET committed
473
        self.controller.execute(WORKER1, 1, CONFIGURATION1)
Philip ABBET's avatar
Philip ABBET committed
474
        self.controller.execute(WORKER1, 2, config)
Philip ABBET's avatar
Philip ABBET committed
475

Philip ABBET's avatar
Philip ABBET committed
476
477
        message = self._wait()
        self._check_done(message, WORKER1, 1)
Philip ABBET's avatar
Philip ABBET committed
478

Philip ABBET's avatar
Philip ABBET committed
479
480
481
        message = self._wait()
        self._check_done(message, WORKER1, 2)

Philip ABBET's avatar
Philip ABBET committed
482

Philip ABBET's avatar
Philip ABBET committed
483
484
485
486
    def test_reuse(self):
        self.controller.execute(WORKER1, 1, CONFIGURATION1)
        message = self._wait()
        self._check_done(message, WORKER1, 1)
Philip ABBET's avatar
Philip ABBET committed
487

Philip ABBET's avatar
Philip ABBET committed
488
489
490
        self.controller.execute(WORKER1, 2, CONFIGURATION1)
        message = self._wait()
        self._check_done(message, WORKER1, 2)
Philip ABBET's avatar
Philip ABBET committed
491
492


Philip ABBET's avatar
Philip ABBET committed
493
494
495
    def test_cancel(self):
        config = dict(CONFIGURATION1)
        config['algorithm'] = 'user/integers_echo_slow/1'
Philip ABBET's avatar
Philip ABBET committed
496

Philip ABBET's avatar
Philip ABBET committed
497
        self.controller.execute(WORKER1, 1, config)
498
        self.controller.cancel(WORKER1, 1)
Philip ABBET's avatar
Philip ABBET committed
499

Philip ABBET's avatar
Philip ABBET committed
500
        (worker, status, job_id, data) = self._wait()
Philip ABBET's avatar
Philip ABBET committed
501

Philip ABBET's avatar
Philip ABBET committed
502
503
504
505
        self.assertEqual(worker, WORKER1)
        self.assertEqual(status, WorkerController.CANCELLED)
        self.assertEqual(job_id, 1)
        self.assertEqual(len(data), 0)
Philip ABBET's avatar
Philip ABBET committed
506
507


508
509
    def test_error_cancel_unknown_job(self):
        self.controller.cancel(WORKER1, 1)
Philip ABBET's avatar
Philip ABBET committed
510

Philip ABBET's avatar
Philip ABBET committed
511
        (worker, status, job_id, data) = self._wait()
Philip ABBET's avatar
Philip ABBET committed
512

Philip ABBET's avatar
Philip ABBET committed
513
514
515
        self.assertEqual(worker, WORKER1)
        self.assertEqual(status, WorkerController.ERROR)
        self.assertTrue(job_id is None)
Philip ABBET's avatar
Philip ABBET committed
516
        self.assertEqual(data[0], "Unknown job: 1")
Philip ABBET's avatar
Philip ABBET committed
517
518


Philip ABBET's avatar
Philip ABBET committed
519
520
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
521
522
523

class TestOneWorkerDocker(TestOneWorker):

Philip ABBET's avatar
Philip ABBET committed
524
525
526
    def __init__(self, methodName='runTest'):
        super(TestOneWorkerDocker, self).__init__(methodName)
        self.docker = True
Philip ABBET's avatar
Philip ABBET committed
527
528


Philip ABBET's avatar
Philip ABBET committed
529
530
531
    @classmethod
    def setUpClass(cls):
        cls.host = Host(raise_on_errors=False)
Philip ABBET's avatar
Philip ABBET committed
532
533


Philip ABBET's avatar
Philip ABBET committed
534
535
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
536

537
class TestTwoWorkers(TestWorkerBase):
Philip ABBET's avatar
Philip ABBET committed
538

Philip ABBET's avatar
Philip ABBET committed
539
540
    def setUp(self):
        self.tearDown()   # In case another test failed badly during its setUp()
Philip ABBET's avatar
Philip ABBET committed
541

542
        super(TestTwoWorkers, self).setUp()
Philip ABBET's avatar
Philip ABBET committed
543

544
545
546
547
548
        self.start_controller()
        self.start_worker(WORKER1)
        self.start_worker(WORKER2)
        self.wait_for_worker_connection(WORKER1)
        self.wait_for_worker_connection(WORKER2)
Philip ABBET's avatar
Philip ABBET committed
549
550


Philip ABBET's avatar
Philip ABBET committed
551
552
    def _test_success_one_worker(self, worker_name):
        self.controller.execute(worker_name, 1, CONFIGURATION1)
Philip ABBET's avatar
Philip ABBET committed
553

Philip ABBET's avatar
Philip ABBET committed
554
555
556
        message = None
        while message is None:
            message = self.controller.process(100)
Philip ABBET's avatar
Philip ABBET committed
557

Philip ABBET's avatar
Philip ABBET committed
558
        (worker, status, job_id, data) = message
Philip ABBET's avatar
Philip ABBET committed
559

Philip ABBET's avatar
Philip ABBET committed
560
561
562
        self.assertEqual(worker, worker_name)
        self.assertEqual(status, WorkerController.DONE)
        self.assertEqual(job_id, 1)
Philip ABBET's avatar
Philip ABBET committed
563

Philip ABBET's avatar
Philip ABBET committed
564
        result = simplejson.loads(data[0])
Philip ABBET's avatar
Philip ABBET committed
565

Philip ABBET's avatar
Philip ABBET committed
566
        self.assertEqual(result['status'], 0)
Philip ABBET's avatar
Philip ABBET committed
567
568


Philip ABBET's avatar
Philip ABBET committed
569
570
    def test_success_worker1(self):
        self._test_success_one_worker(WORKER1)
Philip ABBET's avatar
Philip ABBET committed
571
572


Philip ABBET's avatar
Philip ABBET committed
573
574
    def test_success_worker2(self):
        self._test_success_one_worker(WORKER2)
Philip ABBET's avatar
Philip ABBET committed
575
576


Philip ABBET's avatar
Philip ABBET committed
577
578
579
    def test_success_both_workers(self):
        def _check(worker, status, job_id, data):
            self.assertEqual(status, WorkerController.DONE)
Philip ABBET's avatar
Philip ABBET committed
580

Philip ABBET's avatar
Philip ABBET committed
581
582
583
584
585
            if worker == WORKER1:
                self.assertEqual(job_id, 1)
            else:
                self.assertEqual(worker, WORKER2)
                self.assertEqual(job_id, 2)
Philip ABBET's avatar
Philip ABBET committed
586

Philip ABBET's avatar
Philip ABBET committed
587
588
            result = simplejson.loads(data[0])
            self.assertEqual(result['status'], 0)
Philip ABBET's avatar
Philip ABBET committed
589
590


Philip ABBET's avatar
Philip ABBET committed
591
592
        self.controller.execute(WORKER1, 1, CONFIGURATION1)
        self.controller.execute(WORKER2, 2, CONFIGURATION2)
Philip ABBET's avatar
Philip ABBET committed
593

Philip ABBET's avatar
Philip ABBET committed
594
595
596
        message = None
        while message is None:
            message = self.controller.process(100)
Philip ABBET's avatar
Philip ABBET committed
597

Philip ABBET's avatar
Philip ABBET committed
598
599
        (worker1, status, job_id, data) = message
        _check(worker1, status, job_id, data)
Philip ABBET's avatar
Philip ABBET committed
600

Philip ABBET's avatar
Philip ABBET committed
601
602
603
        message = None
        while message is None:
            message = self.controller.process(100)
Philip ABBET's avatar
Philip ABBET committed
604

Philip ABBET's avatar
Philip ABBET committed
605
606
        (worker2, status, job_id, data) = message
        _check(worker2, status, job_id, data)
Philip ABBET's avatar
Philip ABBET committed
607

Philip ABBET's avatar
Philip ABBET committed
608
        self.assertNotEqual(worker1, worker2)
Philip ABBET's avatar
Philip ABBET committed
609
610


Philip ABBET's avatar
Philip ABBET committed
611
612
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
613
614
615

class TestTwoWorkersDocker(TestTwoWorkers):

Philip ABBET's avatar
Philip ABBET committed
616
617
618
    def __init__(self, methodName='runTest'):
        super(TestTwoWorkersDocker, self).__init__(methodName)
        self.docker = True
Philip ABBET's avatar
Philip ABBET committed
619
620


Philip ABBET's avatar
Philip ABBET committed
621
622
623
    @classmethod
    def setUpClass(cls):
        cls.host = Host(raise_on_errors=False)