src/tests/scheduler_event_call_tests.cpp - mesos - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include <string>
 #include <vector>

 #include <gmock/gmock.h>

 #include <mesos/scheduler.hpp>

 #include <mesos/scheduler/scheduler.hpp>

 #include <process/clock.hpp>
 #include <process/future.hpp>
 #include <process/gmock.hpp>
 #include <process/gtest.hpp>
 #include <process/message.hpp>
 #include <process/owned.hpp>
 #include <process/pid.hpp>

 #include <stout/nothing.hpp>
 #include <stout/try.hpp>

 #include "master/master.hpp"

 #include "master/detector/standalone.hpp"

 #include "tests/containerizer.hpp"
 #include "tests/mesos.hpp"

 using mesos::internal::master::Master;
 using mesos::internal::slave::Slave;

 using mesos::master::detector::MasterDetector;
 using mesos::master::detector::StandaloneMasterDetector;

 using mesos::scheduler::Event;

 using process::Clock;
 using process::Future;
 using process::Message;
 using process::Owned;
 using process::PID;
 using process::UPID;

 using std::string;
 using std::vector;

 using testing::_;
 using testing::AtMost;
 using testing::Eq;

 namespace mesos {
 namespace internal {
 namespace tests {


 // These tests intercept master messages and manually
 // post Event messages to the driver.
 class SchedulerDriverEventTest : public MesosTest {};


 // Ensures that the driver can handle the SUBSCRIBED event.
 TEST_F(SchedulerDriverEventTest, Subscribed)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
   frameworkInfo.set_failover_timeout(Weeks(2).secs());

   // Make sure the initial registration calls 'registered'.
   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

   // Intercept the registration message, send a SUBSCRIBED instead.
   Future<Message> frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   // Ensure that there will be no (re-)registration retries
   // from the scheduler driver.
   Clock::pause();

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   FrameworkRegisteredMessage message;
   ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

   FrameworkID frameworkId = message.framework_id();
   frameworkInfo.mutable_id()->CopyFrom(frameworkId);

   Event event;
   event.set_type(Event::SUBSCRIBED);
   event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

   Future<Nothing> registered;
   EXPECT_CALL(sched, registered(&driver, frameworkId, _))
     .WillOnce(FutureSatisfy(&registered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(registered);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the SUBSCRIBED event
 // after a disconnection with the master.
 TEST_F(SchedulerDriverEventTest, SubscribedDisconnection)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
   frameworkInfo.set_failover_timeout(Weeks(2).secs());

   // Make sure the initial registration calls 'registered'.
   MockScheduler sched;
   StandaloneMasterDetector detector(master.get()->pid);
   TestingMesosSchedulerDriver driver(&sched, &detector, frameworkInfo);

   // Intercept the registration message, send a SUBSCRIBED instead.
   Future<Message> frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   // Ensure that there will be no (re-)registration retries
   // from the scheduler driver.
   Clock::pause();

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   FrameworkRegisteredMessage message;
   ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

   FrameworkID frameworkId = message.framework_id();
   frameworkInfo.mutable_id()->CopyFrom(frameworkId);

   Event event;
   event.set_type(Event::SUBSCRIBED);
   event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

   Future<Nothing> registered;
   EXPECT_CALL(sched, registered(&driver, frameworkId, _))
     .WillOnce(FutureSatisfy(&registered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(registered);

   // Simulate a disconnection and expect a 'reregistered' call.
   EXPECT_CALL(sched, disconnected(&driver));

   Future<Message> frameworkReregisteredMessage =
     DROP_MESSAGE(Eq(FrameworkReregisteredMessage().GetTypeName()), _, _);

   detector.appoint(master.get()->pid);

   AWAIT_READY(frameworkReregisteredMessage);

   Future<Nothing> reregistered;
   EXPECT_CALL(sched, reregistered(&driver, _))
     .WillOnce(FutureSatisfy(&reregistered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(reregistered);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the SUBSCRIBED event
 // after a master failover.
 TEST_F(SchedulerDriverEventTest, SubscribedMasterFailover)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
   frameworkInfo.set_failover_timeout(Weeks(2).secs());

   // Make sure the initial registration calls 'registered'.
   MockScheduler sched;
   StandaloneMasterDetector detector(master.get()->pid);
   TestingMesosSchedulerDriver driver(&sched, &detector, frameworkInfo);

   // Intercept the registration message, send a SUBSCRIBED instead.
   Future<Message> frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   // Ensure that there will be no (re-)registration retries
   // from the scheduler driver.
   Clock::pause();

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   FrameworkRegisteredMessage message;
   ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

   FrameworkID frameworkId = message.framework_id();
   frameworkInfo.mutable_id()->CopyFrom(frameworkId);

   Event event;
   event.set_type(Event::SUBSCRIBED);
   event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

   Future<Nothing> registered;
   EXPECT_CALL(sched, registered(&driver, frameworkId, _))
     .WillOnce(FutureSatisfy(&registered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(registered);

   EXPECT_CALL(sched, disconnected(&driver));

   // Fail over the master and expect a 'reregistered' call.
   // Note that the master sends a registered message for
   // this case (see MESOS-786).
   master->reset();
   master = StartMaster();
   ASSERT_SOME(master);

   frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   detector.appoint(master.get()->pid);

   AWAIT_READY(frameworkRegisteredMessage);

   Future<Nothing> reregistered;
   EXPECT_CALL(sched, reregistered(&driver, _))
     .WillOnce(FutureSatisfy(&reregistered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(reregistered);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the SUBSCRIBED event
 // after a scheduler failover.
 TEST_F(SchedulerDriverEventTest, SubscribedSchedulerFailover)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
   frameworkInfo.set_failover_timeout(Weeks(2).secs());

   // Make sure the initial registration calls 'registered'.
   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

   // Intercept the registration message, send a SUBSCRIBED instead.
   Future<Message> frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   // Ensure that there will be no (re-)registration retries
   // from the scheduler driver.
   Clock::pause();

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   FrameworkRegisteredMessage message;
   ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

   FrameworkID frameworkId = message.framework_id();
   frameworkInfo.mutable_id()->CopyFrom(frameworkId);

   Event event;
   event.set_type(Event::SUBSCRIBED);
   event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

   Future<Nothing> registered;
   EXPECT_CALL(sched, registered(&driver, frameworkId, _))
     .WillOnce(FutureSatisfy(&registered));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(registered);

   // Fail over the scheduler and expect a 'registered' call.
   driver.stop(true);

   MockScheduler sched2;
   MesosSchedulerDriver driver2(
       &sched2, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

   frameworkRegisteredMessage =
     DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver2.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid2 = frameworkRegisteredMessage->to;

   Future<Nothing> registered2;
   EXPECT_CALL(sched2, registered(&driver2, frameworkId, _))
     .WillOnce(FutureSatisfy(&registered2));

   process::post(master.get()->pid, frameworkPid2, event);

   AWAIT_READY(registered2);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle an OFFERS event.
 // Note that this includes the ability to bypass the
 // master when sending framework messages.
 TEST_F(SchedulerDriverEventTest, Offers)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver schedDriver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&schedDriver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   schedDriver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   // Start a slave and capture the offers.
   Future<ResourceOffersMessage> resourceOffersMessage =
     DROP_PROTOBUF(ResourceOffersMessage(), _, _);

   MockExecutor exec(DEFAULT_EXECUTOR_ID);
   TestContainerizer containerizer(&exec);

   Owned<MasterDetector> detector = master.get()->createDetector();

   Try<Owned<cluster::Slave>> slave = StartSlave(detector.get(), &containerizer);
   ASSERT_SOME(slave);

   AWAIT_READY(resourceOffersMessage);

   google::protobuf::RepeatedPtrField<Offer> offers =
     resourceOffersMessage->offers();

   ASSERT_EQ(1, offers.size());

   // Ignore future offer messages.
   DROP_PROTOBUFS(ResourceOffersMessage(), _, _);

   // Send the offers event and expect a 'resourceOffers' call.
   Event event;
   event.set_type(Event::OFFERS);
   event.mutable_offers()->mutable_offers()->CopyFrom(offers);

   Future<Nothing> resourceOffers;
   EXPECT_CALL(sched, resourceOffers(&schedDriver, _))
     .WillOnce(FutureSatisfy(&resourceOffers));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(resourceOffers);

   // To test that the framework -> executor messages are
   // sent directly to the slave, launch a task and send
   // the executor a message.
   EXPECT_CALL(exec, registered(_, _, _, _));

   EXPECT_CALL(exec, launchTask(_, _))
     .WillOnce(SendStatusUpdateFromTask(TASK_RUNNING));

   Future<TaskStatus> status;
   EXPECT_CALL(sched, statusUpdate(&schedDriver, _))
     .WillOnce(FutureArg<1>(&status));

   TaskInfo task = createTask(offers.Get(0), "", DEFAULT_EXECUTOR_ID);

   schedDriver.launchTasks(offers.Get(0).id(), {task});

   AWAIT_READY(status);
   EXPECT_EQ(TASK_RUNNING, status->state());

   // This message should skip the master!
   Future<FrameworkToExecutorMessage> frameworkToExecutorMessage =
     FUTURE_PROTOBUF(
         FrameworkToExecutorMessage(),
         frameworkPid,
         slave.get()->pid);

   Future<string> data;
   EXPECT_CALL(exec, frameworkMessage(_, _))
     .WillOnce(FutureArg<1>(&data));

   schedDriver.sendFrameworkMessage(
       DEFAULT_EXECUTOR_ID, offers.Get(0).slave_id(), "hello");

   AWAIT_READY(frameworkToExecutorMessage);
   AWAIT_EXPECT_EQ("hello", data);

   EXPECT_CALL(exec, shutdown(_))
     .Times(AtMost(1));

   schedDriver.stop();
   schedDriver.join();
 }


 // Ensures that the driver can handle the RESCIND event.
 TEST_F(SchedulerDriverEventTest, Rescind)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&driver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   Event event;
   event.set_type(Event::RESCIND);
   event.mutable_rescind()->mutable_offer_id()->set_value("O");

   Future<Nothing> offerRescinded;
   EXPECT_CALL(sched, offerRescinded(&driver, event.rescind().offer_id()))
     .WillOnce(FutureSatisfy(&offerRescinded));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(offerRescinded);

   driver.stop();
   driver.join();
 }


 // Ensures the scheduler driver can handle the UPDATE event.
 TEST_F(SchedulerDriverEventTest, Update)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&driver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   FrameworkRegisteredMessage message;
   ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

   FrameworkID frameworkId = message.framework_id();

   SlaveID slaveId;
   slaveId.set_value("S");

   TaskID taskId;
   taskId.set_value("T");

   ExecutorID executorId;
   executorId.set_value("E");

   // Generate an update that needs no acknowledgement.
   Event event;
   event.set_type(Event::UPDATE);
   event.mutable_update()->mutable_status()->CopyFrom(
       protobuf::createStatusUpdate(
           frameworkId,
           slaveId,
           taskId,
           TASK_RUNNING,
           TaskStatus::SOURCE_MASTER,
           None(),
           "message",
           None(),
           executorId).status());

   Future<Nothing> statusUpdate;
   Future<Nothing> statusUpdate2;
   EXPECT_CALL(sched, statusUpdate(&driver, event.update().status()))
     .WillOnce(FutureSatisfy(&statusUpdate))
     .WillOnce(FutureSatisfy(&statusUpdate2));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(statusUpdate);

   // Generate an update that requires acknowledgement.
   event.mutable_update()->mutable_status()->set_uuid(UUID::random().toBytes());

   Future<mesos::scheduler::Call> acknowledgement = DROP_CALL(
       mesos::scheduler::Call(), mesos::scheduler::Call::ACKNOWLEDGE, _, _);

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(statusUpdate2);
   AWAIT_READY(acknowledgement);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the MESSAGE event.
 TEST_F(SchedulerDriverEventTest, Message)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&driver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   Event event;
   event.set_type(Event::MESSAGE);
   event.mutable_message()->mutable_slave_id()->set_value("S");
   event.mutable_message()->mutable_executor_id()->set_value("E");
   event.mutable_message()->set_data("data");

   Future<Nothing> frameworkMessage;
   EXPECT_CALL(sched, frameworkMessage(
       &driver,
       event.message().executor_id(),
       event.message().slave_id(),
       event.message().data()))
     .WillOnce(FutureSatisfy(&frameworkMessage));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(frameworkMessage);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the FAILURE event.
 TEST_F(SchedulerDriverEventTest, Failure)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&driver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   // Send a failure for an executor, which should trigger executorLost callback.
   SlaveID slaveId;
   slaveId.set_value("S");

   ExecutorID executorId = DEFAULT_EXECUTOR_ID;

   const int32_t status = 255;

   Event event;
   event.set_type(Event::FAILURE);
   event.mutable_failure()->mutable_slave_id()->CopyFrom(slaveId);
   event.mutable_failure()->mutable_executor_id()->CopyFrom(executorId);
   event.mutable_failure()->set_status(status);

   Future<Nothing> executorLost;
   EXPECT_CALL(sched, executorLost(&driver, executorId, slaveId, status))
     .WillOnce(FutureSatisfy(&executorLost));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(executorLost);

   // Now, post a failure for a slave and expect a 'slaveLost'.
   event.mutable_failure()->clear_executor_id();

   Future<Nothing> slaveLost;
   EXPECT_CALL(sched, slaveLost(&driver, slaveId))
     .WillOnce(FutureSatisfy(&slaveLost));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(slaveLost);

   driver.stop();
   driver.join();
 }


 // Ensures that the driver can handle the ERROR event.
 TEST_F(SchedulerDriverEventTest, Error)
 {
   Try<Owned<cluster::Master>> master = StartMaster();
   ASSERT_SOME(master);

   MockScheduler sched;
   MesosSchedulerDriver driver(
       &sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

   EXPECT_CALL(sched, registered(&driver, _, _));

   Future<Message> frameworkRegisteredMessage =
     FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

   driver.start();

   AWAIT_READY(frameworkRegisteredMessage);
   UPID frameworkPid = frameworkRegisteredMessage->to;

   Event event;
   event.set_type(Event::ERROR);
   event.mutable_error()->set_message("error message");

   Future<Nothing> error;
   EXPECT_CALL(sched, error(&driver, event.error().message()))
     .WillOnce(FutureSatisfy(&error));

   process::post(master.get()->pid, frameworkPid, event);

   AWAIT_READY(error);

   driver.stop();
   driver.join();
 }

 } // namespace tests {
 } // namespace internal {
 } // namespace mesos {
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include <string>
	#include <vector>

	#include <gmock/gmock.h>

	#include <mesos/scheduler.hpp>

	#include <mesos/scheduler/scheduler.hpp>

	#include <process/clock.hpp>
	#include <process/future.hpp>
	#include <process/gmock.hpp>
	#include <process/gtest.hpp>
	#include <process/message.hpp>
	#include <process/owned.hpp>
	#include <process/pid.hpp>

	#include <stout/nothing.hpp>
	#include <stout/try.hpp>

	#include "master/master.hpp"

	#include "master/detector/standalone.hpp"

	#include "tests/containerizer.hpp"
	#include "tests/mesos.hpp"

	using mesos::internal::master::Master;
	using mesos::internal::slave::Slave;

	using mesos::master::detector::MasterDetector;
	using mesos::master::detector::StandaloneMasterDetector;

	using mesos::scheduler::Event;

	using process::Clock;
	using process::Future;
	using process::Message;
	using process::Owned;
	using process::PID;
	using process::UPID;

	using std::string;
	using std::vector;

	using testing::_;
	using testing::AtMost;
	using testing::Eq;

	namespace mesos {
	namespace internal {
	namespace tests {


	// These tests intercept master messages and manually
	// post Event messages to the driver.
	class SchedulerDriverEventTest : public MesosTest {};


	// Ensures that the driver can handle the SUBSCRIBED event.
	TEST_F(SchedulerDriverEventTest, Subscribed)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
	frameworkInfo.set_failover_timeout(Weeks(2).secs());

	// Make sure the initial registration calls 'registered'.
	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

	// Intercept the registration message, send a SUBSCRIBED instead.
	Future<Message> frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	// Ensure that there will be no (re-)registration retries
	// from the scheduler driver.
	Clock::pause();

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	FrameworkRegisteredMessage message;
	ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

	FrameworkID frameworkId = message.framework_id();
	frameworkInfo.mutable_id()->CopyFrom(frameworkId);

	Event event;
	event.set_type(Event::SUBSCRIBED);
	event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

	Future<Nothing> registered;
	EXPECT_CALL(sched, registered(&driver, frameworkId, _))
	.WillOnce(FutureSatisfy(&registered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(registered);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the SUBSCRIBED event
	// after a disconnection with the master.
	TEST_F(SchedulerDriverEventTest, SubscribedDisconnection)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
	frameworkInfo.set_failover_timeout(Weeks(2).secs());

	// Make sure the initial registration calls 'registered'.
	MockScheduler sched;
	StandaloneMasterDetector detector(master.get()->pid);
	TestingMesosSchedulerDriver driver(&sched, &detector, frameworkInfo);

	// Intercept the registration message, send a SUBSCRIBED instead.
	Future<Message> frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	// Ensure that there will be no (re-)registration retries
	// from the scheduler driver.
	Clock::pause();

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	FrameworkRegisteredMessage message;
	ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

	FrameworkID frameworkId = message.framework_id();
	frameworkInfo.mutable_id()->CopyFrom(frameworkId);

	Event event;
	event.set_type(Event::SUBSCRIBED);
	event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

	Future<Nothing> registered;
	EXPECT_CALL(sched, registered(&driver, frameworkId, _))
	.WillOnce(FutureSatisfy(&registered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(registered);

	// Simulate a disconnection and expect a 'reregistered' call.
	EXPECT_CALL(sched, disconnected(&driver));

	Future<Message> frameworkReregisteredMessage =
	DROP_MESSAGE(Eq(FrameworkReregisteredMessage().GetTypeName()), _, _);

	detector.appoint(master.get()->pid);

	AWAIT_READY(frameworkReregisteredMessage);

	Future<Nothing> reregistered;
	EXPECT_CALL(sched, reregistered(&driver, _))
	.WillOnce(FutureSatisfy(&reregistered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(reregistered);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the SUBSCRIBED event
	// after a master failover.
	TEST_F(SchedulerDriverEventTest, SubscribedMasterFailover)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
	frameworkInfo.set_failover_timeout(Weeks(2).secs());

	// Make sure the initial registration calls 'registered'.
	MockScheduler sched;
	StandaloneMasterDetector detector(master.get()->pid);
	TestingMesosSchedulerDriver driver(&sched, &detector, frameworkInfo);

	// Intercept the registration message, send a SUBSCRIBED instead.
	Future<Message> frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	// Ensure that there will be no (re-)registration retries
	// from the scheduler driver.
	Clock::pause();

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	FrameworkRegisteredMessage message;
	ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

	FrameworkID frameworkId = message.framework_id();
	frameworkInfo.mutable_id()->CopyFrom(frameworkId);

	Event event;
	event.set_type(Event::SUBSCRIBED);
	event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

	Future<Nothing> registered;
	EXPECT_CALL(sched, registered(&driver, frameworkId, _))
	.WillOnce(FutureSatisfy(&registered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(registered);

	EXPECT_CALL(sched, disconnected(&driver));

	// Fail over the master and expect a 'reregistered' call.
	// Note that the master sends a registered message for
	// this case (see MESOS-786).
	master->reset();
	master = StartMaster();
	ASSERT_SOME(master);

	frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	detector.appoint(master.get()->pid);

	AWAIT_READY(frameworkRegisteredMessage);

	Future<Nothing> reregistered;
	EXPECT_CALL(sched, reregistered(&driver, _))
	.WillOnce(FutureSatisfy(&reregistered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(reregistered);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the SUBSCRIBED event
	// after a scheduler failover.
	TEST_F(SchedulerDriverEventTest, SubscribedSchedulerFailover)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	FrameworkInfo frameworkInfo = DEFAULT_FRAMEWORK_INFO;
	frameworkInfo.set_failover_timeout(Weeks(2).secs());

	// Make sure the initial registration calls 'registered'.
	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

	// Intercept the registration message, send a SUBSCRIBED instead.
	Future<Message> frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	// Ensure that there will be no (re-)registration retries
	// from the scheduler driver.
	Clock::pause();

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	FrameworkRegisteredMessage message;
	ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

	FrameworkID frameworkId = message.framework_id();
	frameworkInfo.mutable_id()->CopyFrom(frameworkId);

	Event event;
	event.set_type(Event::SUBSCRIBED);
	event.mutable_subscribed()->mutable_framework_id()->CopyFrom(frameworkId);

	Future<Nothing> registered;
	EXPECT_CALL(sched, registered(&driver, frameworkId, _))
	.WillOnce(FutureSatisfy(&registered));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(registered);

	// Fail over the scheduler and expect a 'registered' call.
	driver.stop(true);

	MockScheduler sched2;
	MesosSchedulerDriver driver2(
	&sched2, frameworkInfo, master.get()->pid, DEFAULT_CREDENTIAL);

	frameworkRegisteredMessage =
	DROP_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver2.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid2 = frameworkRegisteredMessage->to;

	Future<Nothing> registered2;
	EXPECT_CALL(sched2, registered(&driver2, frameworkId, _))
	.WillOnce(FutureSatisfy(&registered2));

	process::post(master.get()->pid, frameworkPid2, event);

	AWAIT_READY(registered2);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle an OFFERS event.
	// Note that this includes the ability to bypass the
	// master when sending framework messages.
	TEST_F(SchedulerDriverEventTest, Offers)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver schedDriver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&schedDriver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	schedDriver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	// Start a slave and capture the offers.
	Future<ResourceOffersMessage> resourceOffersMessage =
	DROP_PROTOBUF(ResourceOffersMessage(), _, _);

	MockExecutor exec(DEFAULT_EXECUTOR_ID);
	TestContainerizer containerizer(&exec);

	Owned<MasterDetector> detector = master.get()->createDetector();

	Try<Owned<cluster::Slave>> slave = StartSlave(detector.get(), &containerizer);
	ASSERT_SOME(slave);

	AWAIT_READY(resourceOffersMessage);

	google::protobuf::RepeatedPtrField<Offer> offers =
	resourceOffersMessage->offers();

	ASSERT_EQ(1, offers.size());

	// Ignore future offer messages.
	DROP_PROTOBUFS(ResourceOffersMessage(), _, _);

	// Send the offers event and expect a 'resourceOffers' call.
	Event event;
	event.set_type(Event::OFFERS);
	event.mutable_offers()->mutable_offers()->CopyFrom(offers);

	Future<Nothing> resourceOffers;
	EXPECT_CALL(sched, resourceOffers(&schedDriver, _))
	.WillOnce(FutureSatisfy(&resourceOffers));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(resourceOffers);

	// To test that the framework -> executor messages are
	// sent directly to the slave, launch a task and send
	// the executor a message.
	EXPECT_CALL(exec, registered(_, _, _, _));

	EXPECT_CALL(exec, launchTask(_, _))
	.WillOnce(SendStatusUpdateFromTask(TASK_RUNNING));

	Future<TaskStatus> status;
	EXPECT_CALL(sched, statusUpdate(&schedDriver, _))
	.WillOnce(FutureArg<1>(&status));

	TaskInfo task = createTask(offers.Get(0), "", DEFAULT_EXECUTOR_ID);

	schedDriver.launchTasks(offers.Get(0).id(), {task});

	AWAIT_READY(status);
	EXPECT_EQ(TASK_RUNNING, status->state());

	// This message should skip the master!
	Future<FrameworkToExecutorMessage> frameworkToExecutorMessage =
	FUTURE_PROTOBUF(
	FrameworkToExecutorMessage(),
	frameworkPid,
	slave.get()->pid);

	Future<string> data;
	EXPECT_CALL(exec, frameworkMessage(_, _))
	.WillOnce(FutureArg<1>(&data));

	schedDriver.sendFrameworkMessage(
	DEFAULT_EXECUTOR_ID, offers.Get(0).slave_id(), "hello");

	AWAIT_READY(frameworkToExecutorMessage);
	AWAIT_EXPECT_EQ("hello", data);

	EXPECT_CALL(exec, shutdown(_))
	.Times(AtMost(1));

	schedDriver.stop();
	schedDriver.join();
	}


	// Ensures that the driver can handle the RESCIND event.
	TEST_F(SchedulerDriverEventTest, Rescind)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&driver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	Event event;
	event.set_type(Event::RESCIND);
	event.mutable_rescind()->mutable_offer_id()->set_value("O");

	Future<Nothing> offerRescinded;
	EXPECT_CALL(sched, offerRescinded(&driver, event.rescind().offer_id()))
	.WillOnce(FutureSatisfy(&offerRescinded));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(offerRescinded);

	driver.stop();
	driver.join();
	}


	// Ensures the scheduler driver can handle the UPDATE event.
	TEST_F(SchedulerDriverEventTest, Update)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&driver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	FrameworkRegisteredMessage message;
	ASSERT_TRUE(message.ParseFromString(frameworkRegisteredMessage->body));

	FrameworkID frameworkId = message.framework_id();

	SlaveID slaveId;
	slaveId.set_value("S");

	TaskID taskId;
	taskId.set_value("T");

	ExecutorID executorId;
	executorId.set_value("E");

	// Generate an update that needs no acknowledgement.
	Event event;
	event.set_type(Event::UPDATE);
	event.mutable_update()->mutable_status()->CopyFrom(
	protobuf::createStatusUpdate(
	frameworkId,
	slaveId,
	taskId,
	TASK_RUNNING,
	TaskStatus::SOURCE_MASTER,
	None(),
	"message",
	None(),
	executorId).status());

	Future<Nothing> statusUpdate;
	Future<Nothing> statusUpdate2;
	EXPECT_CALL(sched, statusUpdate(&driver, event.update().status()))
	.WillOnce(FutureSatisfy(&statusUpdate))
	.WillOnce(FutureSatisfy(&statusUpdate2));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(statusUpdate);

	// Generate an update that requires acknowledgement.
	event.mutable_update()->mutable_status()->set_uuid(UUID::random().toBytes());

	Future<mesos::scheduler::Call> acknowledgement = DROP_CALL(
	mesos::scheduler::Call(), mesos::scheduler::Call::ACKNOWLEDGE, _, _);

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(statusUpdate2);
	AWAIT_READY(acknowledgement);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the MESSAGE event.
	TEST_F(SchedulerDriverEventTest, Message)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&driver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	Event event;
	event.set_type(Event::MESSAGE);
	event.mutable_message()->mutable_slave_id()->set_value("S");
	event.mutable_message()->mutable_executor_id()->set_value("E");
	event.mutable_message()->set_data("data");

	Future<Nothing> frameworkMessage;
	EXPECT_CALL(sched, frameworkMessage(
	&driver,
	event.message().executor_id(),
	event.message().slave_id(),
	event.message().data()))
	.WillOnce(FutureSatisfy(&frameworkMessage));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(frameworkMessage);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the FAILURE event.
	TEST_F(SchedulerDriverEventTest, Failure)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&driver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	// Send a failure for an executor, which should trigger executorLost callback.
	SlaveID slaveId;
	slaveId.set_value("S");

	ExecutorID executorId = DEFAULT_EXECUTOR_ID;

	const int32_t status = 255;

	Event event;
	event.set_type(Event::FAILURE);
	event.mutable_failure()->mutable_slave_id()->CopyFrom(slaveId);
	event.mutable_failure()->mutable_executor_id()->CopyFrom(executorId);
	event.mutable_failure()->set_status(status);

	Future<Nothing> executorLost;
	EXPECT_CALL(sched, executorLost(&driver, executorId, slaveId, status))
	.WillOnce(FutureSatisfy(&executorLost));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(executorLost);

	// Now, post a failure for a slave and expect a 'slaveLost'.
	event.mutable_failure()->clear_executor_id();

	Future<Nothing> slaveLost;
	EXPECT_CALL(sched, slaveLost(&driver, slaveId))
	.WillOnce(FutureSatisfy(&slaveLost));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(slaveLost);

	driver.stop();
	driver.join();
	}


	// Ensures that the driver can handle the ERROR event.
	TEST_F(SchedulerDriverEventTest, Error)
	{
	Try<Owned<cluster::Master>> master = StartMaster();
	ASSERT_SOME(master);

	MockScheduler sched;
	MesosSchedulerDriver driver(
	&sched, DEFAULT_FRAMEWORK_INFO, master.get()->pid, DEFAULT_CREDENTIAL);

	EXPECT_CALL(sched, registered(&driver, _, _));

	Future<Message> frameworkRegisteredMessage =
	FUTURE_MESSAGE(Eq(FrameworkRegisteredMessage().GetTypeName()), _, _);

	driver.start();

	AWAIT_READY(frameworkRegisteredMessage);
	UPID frameworkPid = frameworkRegisteredMessage->to;

	Event event;
	event.set_type(Event::ERROR);
	event.mutable_error()->set_message("error message");

	Future<Nothing> error;
	EXPECT_CALL(sched, error(&driver, event.error().message()))
	.WillOnce(FutureSatisfy(&error));

	process::post(master.get()->pid, frameworkPid, event);

	AWAIT_READY(error);

	driver.stop();
	driver.join();
	}

	} // namespace tests {
	} // namespace internal {
	} // namespace mesos {