def test_accept_offers(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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_gen_request(mocker): mock_addr = 'mock_addr:1234' sched = mocker.Mock() framework = {} 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_on_error(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) msg = 'error message' event = {'type': 'ERROR', 'error': {'message': msg}} driver.on_event(event) sched.error.assert_called_once_with(driver, msg)
def test_on_rescind(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) offer_id = 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 test_teardown(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) driver._send = mocker.Mock() driver._teardown() driver._send.assert_called_once_with({ 'type': 'TEARDOWN', 'framework_id': { 'value': ID }, })
def test_on_offers(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) 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_send(mocker): sched = mocker.Mock() framework = mocker.Mock() master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) resp = mocker.Mock(status=200, read=mocker.Mock(return_value='{}')) conn = mocker.Mock(getresponse=mocker.Mock(return_value=resp)) driver._get_conn = mocker.Mock(return_value=conn) assert driver._send({}) == {} driver._get_conn.assert_called_once_with() conn.request.assert_called_once_with( 'POST', '/api/v1/scheduler', body=json.dumps({}).encode('utf-8'), headers={'Content-Type': 'application/json'})
def test_decline_offer(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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 } })
def test_reconcile_tasks(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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_on_failure(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) executor_id = str(uuid.uuid4()) agent_id = str(uuid.uuid4()) status = random.randint(0, 256) event = { 'type': 'FAILURE', 'failure': { 'executor_id': executor_id, 'agent_id': agent_id, 'status': status } } driver.on_event(event) sched.executorLost.assert_called_once_with(driver, executor_id, agent_id, status) event = { 'type': 'FAILURE', 'failure': { 'agent_id': agent_id, } } driver.on_event(event) sched.slaveLost.assert_called_once_with(driver, agent_id)
def test_on_message(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) executor_id = str(uuid.uuid4()) agent_id = str(uuid.uuid4()) message = ''.join( random.choice(string.printable) for _ in range(random.randint(1, 100))) data = b2a_base64(message.encode('utf8')) event = { 'type': 'MESSAGE', 'message': { 'executor_id': executor_id, 'agent_id': agent_id, 'data': data } } driver.on_event(event) sched.frameworkMessage.assert_called_once_with(driver, executor_id, agent_id, data)
def test_request_resources(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': ID} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) driver._send = mocker.Mock() requests = [{ 'agent_id': { 'value': str(uuid.uuid4()) }, 'resources': {} } for _ in range(random.randint(1, 10))] driver.requestResources(requests) driver._send.assert_called_once_with({ 'type': 'REQUEST', 'framework_id': { 'value': ID }, 'request': { 'requests': requests } })
def test_send_framework_message(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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))) 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': b2a_base64(message.encode('utf-8')).rstrip(), } })
def test_acknowledge_status_update(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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 = b2a_base64(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_revive_offers(mocker): ID = str(uuid.uuid4()) sched = mocker.Mock() framework = {'id': 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_on_subscribed(mocker): sched = mocker.Mock() framework = {} master = mocker.Mock() driver = MesosSchedulerDriver(sched, framework, master) driver.version = '1.0.0' driver._master = 'mock_addr:12345' framework_id = {'value': str(uuid.uuid4())} event = { 'type': 'SUBSCRIBED', 'subscribed': { 'framework_id': framework_id } } driver.on_event(event) sched.registered.assert_called_once_with(driver, framework_id, { 'hostname': 'mock_addr', 'port': 12345, 'version': '1.0.0' })