Python pymesos.MesosSchedulerDriver类代码示例

def test_gen_request(mocker):
    mock_addr = 'mock_addr:1234'
    sched = mocker.Mock()
    framework = {
        'failover_timeout': 0
    }
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._master = mock_addr
    req = driver.gen_request()
    parser = HttpParser(0)
    assert len(req) == parser.execute(req, len(req))
    assert parser.is_headers_complete()
    assert parser.get_method() == 'POST'
    assert parser.get_url() == '/api/v1/scheduler'

    assert parser.is_partial_body()
    body = parser.recv_body()
    result = json.loads(body.decode('utf-8'))
    assert result['type'] == 'SUBSCRIBE'
    assert result['subscribe'] == {
        'framework_info': framework
    }

    headers = {k.upper(): v for k, v in parser.get_headers().items()}
    assert headers == {
        'HOST': mock_addr,
        'CONTENT-TYPE': 'application/json',
        'ACCEPT': 'application/json',
        'CONNECTION': 'close',
        'CONTENT-LENGTH': str(len(body))
    }

    assert parser.is_message_complete()

def test_accept_offers(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    offer_ids = [str(uuid.uuid4()) for _ in range(random.randint(1, 10))]
    operations = [{
        'type': 'LAUNCH',
        'launch': {
            'task_infos': [
                {
                    'name': '1',
                },
                {
                    'name': '2',
                }
            ]
        }
    }]
    driver.acceptOffers(offer_ids, operations)
    driver._send.assert_called_once_with({
        'type': 'ACCEPT',
        'framework_id': {
            'value': ID
        },
        'accept': {
            'offer_ids': offer_ids,
            'operations': operations,
        }
    })

def test_acknowledge_status_update(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    agent_id = dict(value=str(uuid.uuid4()))
    task_id = dict(value=str(uuid.uuid4()))
    uid = encode_data(uuid.uuid4().bytes)
    status = {
        'agent_id': agent_id,
        'task_id': task_id,
        'uuid': uid
    }
    driver.acknowledgeStatusUpdate(status)
    driver._send.assert_called_once_with({
        'type': 'ACKNOWLEDGE',
        'framework_id': {
            'value': ID
        },
        'acknowledge': {
            'agent_id': agent_id,
            'task_id': task_id,
            'uuid': uid
        }
    })

def test_reconcile_tasks(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    task_ids = [str(uuid.uuid4()) for _ in range(random.randint(1, 10))]
    tasks = [
        {
            'task_id': {
                'value': id
            }
        }
        for id in task_ids
    ]
    driver.reconcileTasks(tasks)
    driver._send.assert_called_once_with({
        'type': 'RECONCILE',
        'framework_id': {
            'value': ID
        },
        'reconcile': {
            'tasks': tasks
        }
    })

def test_revive_offers(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    driver._stream_id = str(uuid.uuid4())
    driver.reviveOffers()
    driver._send.assert_called_once_with({
        'type': 'REVIVE',
        'framework_id': {
            'value': ID
        },
    })

def test_teardown(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    driver.stream_id = str(uuid.uuid4())
    assert driver.connected
    driver._teardown()
    driver._send.assert_called_once_with({
        'type': 'TEARDOWN',
        'framework_id': {
            'value': ID
        },
    })

def test_on_rescind(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._started = True
    offer_id = {'value': str(uuid.uuid4())}
    event = {
        'type': 'RESCIND',
        'rescind': {
            'offer_id': offer_id
        }
    }
    driver.on_event(event)
    sched.offerRescinded.assert_called_once_with(driver, offer_id)

    def start_driver(self):
        name = '[dpark] ' + \
            os.path.abspath(sys.argv[0]) + ' ' + ' '.join(sys.argv[1:])
        if len(name) > 256:
            name = name[:256] + '...'
        framework = Dict()
        framework.user = getuser()
        if framework.user == 'root':
            raise Exception('dpark is not allowed to run as \'root\'')
        framework.name = name
        framework.hostname = socket.gethostname()
        framework.webui_url = self.options.webui_url

        self.driver = MesosSchedulerDriver(
            self, framework, self.master, use_addict=True
        )
        self.driver.start()
        logger.debug('Mesos Scheudler driver started')

        self.started = True
        self.last_finish_time = time.time()

        def check():
            while self.started:
                now = time.time()
                if (not self.activeJobs and
                        now - self.last_finish_time > MAX_IDLE_TIME):
                    logger.info('stop mesos scheduler after %d seconds idle',
                                now - self.last_finish_time)
                    self.stop()
                    break
                time.sleep(1)

        spawn(check)

    def start(self):

        def readable(fd):
            return bool(select.select([fd], [], [], 0.1)[0])

        lfd = socket.socket()
        try:
            lfd.bind(('', 0))
            self.addr = '%s:%s' % (socket.gethostname(), lfd.getsockname()[1])
            lfd.listen(10)
            framework = mesos_pb2.FrameworkInfo()
            framework.user = getpass.getuser()
            framework.name = self.name
            framework.hostname = socket.gethostname()
            self.driver = MesosSchedulerDriver(self, framework, self.master)
            self.driver.start()
            while any((not task.initalized for task in self.tasks)):
                if readable(lfd):
                    c, _ = lfd.accept()
                    if readable(c):
                        mesos_task_id, addr = recv(c)
                        assert isinstance(mesos_task_id, int)
                        task = self.tasks[mesos_task_id]
                        task.addr = addr
                        task.connection = c
                        task.initalized = True
                    else:
                        c.close()
            return self._start_tf_cluster()
        except Exception:
            self.stop()
            raise
        finally:
            lfd.close()

def test_on_error(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._started = True
    msg = 'error message'
    event = {
        'type': 'ERROR',
        'error': {
            'message': msg
        }
    }
    driver.on_event(event)
    sched.error.assert_called_once_with(driver, msg)

def test_send_framework_message(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    executor_id = {'value': str(uuid.uuid4())}
    agent_id = {'value': str(uuid.uuid4())}
    message = ''.join(random.choice(string.printable)
                      for _ in range(random.randint(1, 100)))
    message = encode_data(message.encode('utf-8'))
    driver.sendFrameworkMessage(executor_id, agent_id, message)
    driver._send.assert_not_called()
    driver._stream_id = 'a-stream-id'
    driver.sendFrameworkMessage(executor_id, agent_id, message)
    driver._send.assert_called_once_with({
        'type': 'MESSAGE',
        'framework_id': {
            'value': ID
        },
        'message': {
            'agent_id': agent_id,
            'executor_id': executor_id,
            'data': message,
        }
    })

def test_reconcile_operations(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    op_ids = [str(uuid.uuid4()) for _ in range(random.randint(1, 10))]
    operations = [
        {
            'operation_id': {
                'value': id
            }
        }
        for id in op_ids
    ]
    driver.reconcileOperations(operations)
    driver._send.assert_not_called()
    driver._stream_id = 'a-stream-id'
    driver.reconcileOperations(operations)
    driver._send.assert_called_once_with({
        'type': 'RECONCILE_OPERATIONS',
        'framework_id': {
            'value': ID
        },
        'reconcile_operations': {
            'operations': operations
        }
    })

def test_acknowledge_operation_status_update(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    agent_id = dict(value=str(uuid.uuid4()))
    operation_id = dict(value=str(uuid.uuid4()))
    uid = encode_data(uuid.uuid4().bytes)
    status = {
        'agent_id': agent_id,
        'operation_id': operation_id,
        'uuid': uid
    }
    driver.acknowledgeOperationStatusUpdate(status)
    driver._send.assert_not_called()
    driver._stream_id = 'a-stream-id'
    driver.acknowledgeOperationStatusUpdate(status)
    driver._send.assert_called_once_with({
        'type': 'ACKNOWLEDGE_OPERATION_STATUS',
        'framework_id': {
            'value': ID
        },
        'acknowledge_operation_status': {
            'agent_id': agent_id,
            'operation_id': operation_id,
            'uuid': uid,
        }
    })

def test_on_offers(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._started = True
    offers = [{
        'offer_id': {'value': str(uuid.uuid4())}
    } for _ in range(random.randint(1, 10))]
    event = {
        'type': 'OFFERS',
        'offers': {
            'offers': offers
        }
    }
    driver.on_event(event)
    sched.resourceOffers.assert_called_once_with(driver, offers)

def test_decline_offer(mocker):
    ID = str(uuid.uuid4())
    sched = mocker.Mock()
    framework = {'id': {'value': ID}}
    master = mocker.Mock()
    driver = MesosSchedulerDriver(sched, framework, master)
    driver._send = mocker.Mock()
    offer_ids = [str(uuid.uuid4()) for _ in range(random.randint(1, 10))]
    driver.declineOffer(offer_ids)
    driver._send.assert_called_once_with({
        'type': 'DECLINE',
        'framework_id': {
            'value': ID
        },
        'decline': {
            'offer_ids': offer_ids
        }
    })