def test_scheduler_event_handlers(mocker): sched = mocker.Mock() driver = mocker.Mock() proxy = SchedulerProxy(sched) proxy.executorLost(driver, mesos_pb2.ExecutorInfo(), mesos_pb2.ExecutorID(), 1) proxy.frameworkMessage(driver, mesos_pb2.ExecutorID(), mesos_pb2.SlaveID(), 'message') proxy.offerRescinded(driver, mesos_pb2.OfferID()) proxy.registered(driver, mesos_pb2.FrameworkID(), mesos_pb2.MasterInfo()) proxy.resourceOffers(driver, [mesos_pb2.Offer(), mesos_pb2.Offer()]) proxy.slaveLost(driver, mesos_pb2.SlaveID()) proxy.statusUpdate(driver, mesos_pb2.TaskStatus()) proxy.reregistered(driver, mesos_pb2.MasterInfo()) proxy.error(driver, 'message') proxy.disconnected(driver) sched.on_executor_lost.assert_called_once() sched.on_message.assert_called_once() sched.on_rescinded.assert_called_once() sched.on_registered.assert_called_once() sched.on_offers.assert_called_once() sched.on_slave_lost.assert_called_once() sched.on_update.assert_called_once() sched.on_reregistered.assert_called_once() sched.on_error.assert_called_once() sched.on_disconnected.assert_called_once()
def make_fake_offer(cpu=50000, mem=50000, port_begin=31000, port_end=32000, pool='default'): offer = mesos_pb2.Offer() offer.slave_id.value = "super big slave" cpus_resource = offer.resources.add() cpus_resource.name = "cpus" cpus_resource.scalar.value = cpu mem_resource = offer.resources.add() mem_resource.name = "mem" mem_resource.scalar.value = mem ports_resource = offer.resources.add() ports_resource.name = "ports" ports_range = ports_resource.ranges.range.add() ports_range.begin = port_begin ports_range.end = port_end if pool is not None: pool_attribute = offer.attributes.add() pool_attribute.name = "pool" pool_attribute.text.value = pool return offer
def setUp(self): self._driver = FakeDriver() self._storage = FakeStorage(SequentialThreadingHandler()) self._zk_client = FakeClient(storage=self._storage) self._zk_client.start() self._framework_id = mesos_pb2.FrameworkID() self._framework_id.value = "framework_id_0" self._offer = mesos_pb2.Offer() self._offer.id.value = "offer_id_0" self._offer.framework_id.value = self._framework_id.value self._offer.slave_id.value = "slave_id_0" self._offer.hostname = "localhost" resources = create_resources(cpus=4, mem=512 * 3, ports=set([10000, 10001, 10002])) self._offer.resources.extend(resources) self._framework_user = "******" self._zk_url = "zk://host/mysos/test" self._cluster = MySQLCluster("cluster0", "user", "pass", 3) self._tmpdir = tempfile.mkdtemp() self._state_provider = LocalStateProvider(self._tmpdir) framework_info = mesos_pb2.FrameworkInfo(user=getpass.getuser(), name="mysos", checkpoint=False) self._state = Scheduler(framework_info)
def setup(self, request): self._driver = FakeDriver() self._storage = FakeStorage(SequentialThreadingHandler()) self._zk_client = FakeClient(storage=self._storage) self._zk_client.start() self._offer = mesos_pb2.Offer() self._offer.id.value = "offer_id_0" self._offer.framework_id.value = "framework_id_0" self._offer.slave_id.value = "slave_id_0" self._offer.hostname = "localhost" # Enough memory and ports to fit three tasks. resources = create_resources(cpus=4, mem=512 * 3, ports=set([10000, 10001, 10002])) self._offer.resources.extend(resources) self._framework_user = "******" # Some tests use the default launcher; some don't. self._zk_url = "zk://host/mysos/test" self._cluster = MySQLCluster("cluster0", "user", "pass", 3) # Construct the state provider based on the test parameter. if request.param == LocalStateProvider: tmpdir = tempfile.mkdtemp() self._state_provider = LocalStateProvider(tmpdir) request.addfinalizer(lambda: shutil.rmtree(tmpdir, True) ) # Clean up after ourselves. elif request.param == ZooKeeperStateProvider: self._state_provider = ZooKeeperStateProvider( self._zk_client, "/mysos/test") self._launcher = MySQLClusterLauncher( self._driver, self._cluster, self._state_provider, self._zk_url, self._zk_client, self._framework_user, "./executor.pex", "cmd.sh", Amount(5, Time.SECONDS), "/etc/mysos/admin_keyfile.yml", query_interval=Amount(150, Time.MILLISECONDS)) # Short interval. self._elected = threading.Event() self._launchers = [self._launcher] # See teardown(). request.addfinalizer(self.teardown)
def test_mesos_executor(self, driver): # create task queue, empty result queue, task_cpu and task_memory tasks_queue = Queue() fake_af_task1 = {"key1", "airflow run tutorial"} fake_af_task2 = {"key2", "airflow run tutorial2"} tasks_queue.put(fake_af_task1) tasks_queue.put(fake_af_task2) results_queue = Queue() task_cpu = 2 task_memory = 4 scheduler = AirflowMesosScheduler(tasks_queue, results_queue, task_cpu, task_memory) # Create Offers resources = [] fake_cpu_resource = mesos_pb2.Resource( name='cpus', type=mesos_pb2.Value.SCALAR, ) fake_cpu_resource.scalar.value = task_cpu fake_mem_resource = mesos_pb2.Resource( name='mem', type=mesos_pb2.Value.SCALAR, ) fake_mem_resource.scalar.value = task_memory resources.append(fake_cpu_resource) resources.append(fake_mem_resource) fake_offer = mesos_pb2.Offer( id=self.offer_id, framework_id=self.framework_id, slave_id=self.slave_id, hostname='fake-host', resources=resources ) scheduler.resourceOffers(driver, [fake_offer]) # assertions self.assertTrue(driver.launchTasks.called)