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

316
317
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
        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)


346
347
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
    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')


374
375
376
377
378
379
380
381
382
383
384
385
#----------------------------------------------------------


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
386
387


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

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

Philip ABBET's avatar
Philip ABBET committed
396
        return message
Philip ABBET's avatar
Philip ABBET committed
397
398


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

Philip ABBET's avatar
Philip ABBET committed
402
        (worker, status, job_id, data) = message
Philip ABBET's avatar
Philip ABBET committed
403

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

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

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


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

Philip ABBET's avatar
Philip ABBET committed
416
        message = self._wait()
Philip ABBET's avatar
Philip ABBET committed
417

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


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

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

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

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

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

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


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

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

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

Philip ABBET's avatar
Philip ABBET committed
447
448
449
450
        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
451
452


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

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

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

Philip ABBET's avatar
Philip ABBET committed
461
462
463
464
        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
465
466


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

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

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

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

Philip ABBET's avatar
Philip ABBET committed
480

Philip ABBET's avatar
Philip ABBET committed
481
482
483
484
    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
485

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


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

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

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

Philip ABBET's avatar
Philip ABBET committed
500
501
502
503
        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
504
505


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

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

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


Philip ABBET's avatar
Philip ABBET committed
517
518
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
519
520
521

class TestOneWorkerDocker(TestOneWorker):

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


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


Philip ABBET's avatar
Philip ABBET committed
532
533
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
534

535
class TestTwoWorkers(TestWorkerBase):
Philip ABBET's avatar
Philip ABBET committed
536

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

540
        super(TestTwoWorkers, self).setUp()
Philip ABBET's avatar
Philip ABBET committed
541

542
543
544
545
546
        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
547
548


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

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

Philip ABBET's avatar
Philip ABBET committed
556
        (worker, status, job_id, data) = message
Philip ABBET's avatar
Philip ABBET committed
557

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

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

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


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


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


Philip ABBET's avatar
Philip ABBET committed
575
576
577
    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
578

Philip ABBET's avatar
Philip ABBET committed
579
580
581
582
583
            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
584

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


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

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

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

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

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

Philip ABBET's avatar
Philip ABBET committed
606
        self.assertNotEqual(worker1, worker2)
Philip ABBET's avatar
Philip ABBET committed
607
608


Philip ABBET's avatar
Philip ABBET committed
609
610
#----------------------------------------------------------

Philip ABBET's avatar
Philip ABBET committed
611
612
613

class TestTwoWorkersDocker(TestTwoWorkers):

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


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