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
32

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

import unittest
import simplejson
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


50 51 52
#----------------------------------------------------------


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


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
}


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
}


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()


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


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


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 = []
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)
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

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


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


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


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)