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)