Added a scheduler API test for a valid offer constraints update.

asekretenko · asekretenko · commit f284314ba053 · 2020-08-20T10:39:22.000+02:00
Review: https://reviews.apache.org/r/72744
diff --git a/src/tests/master/update_framework_tests.cpp b/src/tests/master/update_framework_tests.cpp
@@ -193,7 +193,9 @@ class UpdateFrameworkTest : public MesosTest {};
 static Future<APIResult> callUpdateFramework(
     Mesos* mesos,
     const FrameworkInfo& info,
-    const vector<string>& suppressedRoles = {})
+    const vector<string>& suppressedRoles = {},
+    const Option<::mesos::v1::scheduler::OfferConstraints>& offerConstraints =
+      None())
 {
   CHECK(info.has_id());
 
@@ -203,6 +205,12 @@ static Future<APIResult> callUpdateFramework(
   *call.mutable_update_framework()->mutable_framework_info() = info;
   *call.mutable_update_framework()->mutable_suppressed_roles() =
     RepeatedPtrField<string>(suppressedRoles.begin(), suppressedRoles.end());
+
+  if (offerConstraints.isSome()) {
+    *call.mutable_update_framework()->mutable_offer_constraints() =
+      *offerConstraints;
+  }
+
   return mesos->call(call);
 }
 
@@ -725,6 +733,103 @@ TEST_F(UpdateFrameworkTest, RemoveAndUnsuppress)
 }
 
 
+// This test ensures that it is possible to modify offer constraints
+// via the UpdateFramework call.
+TEST_F(UpdateFrameworkTest, OfferConstraints)
+{
+  using ::mesos::v1::scheduler::AttributeConstraint;
+  using ::mesos::v1::scheduler::OfferConstraints;
+
+  mesos::internal::master::Flags masterFlags = CreateMasterFlags();
+  Try<Owned<cluster::Master>> master = StartMaster(masterFlags);
+  ASSERT_SOME(master);
+
+  Owned<MasterDetector> detector = master->get()->createDetector();
+
+  mesos::internal::slave::Flags slaveFlags = CreateSlaveFlags();
+  Try<Owned<cluster::Slave>> slave = StartSlave(detector.get(), slaveFlags);
+  ASSERT_SOME(slave);
+
+  auto scheduler = std::make_shared<MockHTTPScheduler>();
+
+  EXPECT_CALL(*scheduler, connected(_))
+    .WillOnce(Invoke([](Mesos* mesos) {
+      Call call;
+      call.set_type(Call::SUBSCRIBE);
+      *call.mutable_subscribe()->mutable_framework_info() =
+        DEFAULT_FRAMEWORK_INFO;
+
+      AttributeConstraint* constraint =
+        (*call.mutable_subscribe()
+            ->mutable_offer_constraints()
+            ->mutable_role_constraints())[DEFAULT_FRAMEWORK_INFO.roles(0)]
+          .add_groups()
+          ->add_attribute_constraints();
+
+      *constraint->mutable_selector()->mutable_attribute_name() = "foo";
+      *constraint->mutable_predicate()->mutable_exists() =
+        AttributeConstraint::Predicate::Exists();
+
+      mesos->send(call);
+    }));
+
+  EXPECT_CALL(*scheduler, heartbeat(_))
+    .WillRepeatedly(Return()); // Ignore heartbeats.
+
+  Future<Event::Subscribed> subscribed;
+
+  EXPECT_CALL(*scheduler, subscribed(_, _)).WillOnce(FutureArg<1>(&subscribed));
+
+  // Expect that the framework gets no offers.
+  EXPECT_CALL(*scheduler, offers(_, _)).Times(AtMost(0));
+
+  TestMesos mesos(master->get()->pid, ContentType::PROTOBUF, scheduler);
+
+  AWAIT_READY(subscribed);
+
+  // Trigger allocation to ensure that the agent is not offered before changing
+  // offer constraints.
+  Clock::pause();
+  Clock::settle();
+  Clock::advance(masterFlags.allocation_interval);
+  Clock::settle();
+
+  // Expect an offer after constraints change.
+  Future<Event::Offers> offers;
+  EXPECT_CALL(*scheduler, offers(_, _)).WillOnce(FutureArg<1>(&offers));
+
+  // Change constraint to `NotExists` so that the agent will now be offered to
+  // the framework.
+  {
+    FrameworkInfo framework = DEFAULT_FRAMEWORK_INFO;
+    *framework.mutable_id() = subscribed->framework_id();
+
+    OfferConstraints constraints;
+    AttributeConstraint* constraint =
+      (*constraints.mutable_role_constraints())[framework.roles(0)]
+        .add_groups()
+        ->add_attribute_constraints();
+
+    *constraint->mutable_selector()->mutable_attribute_name() = "foo";
+    *constraint->mutable_predicate()->mutable_not_exists() =
+      AttributeConstraint::Predicate::NotExists();
+
+    AWAIT_READY(callUpdateFramework(&mesos, framework, {}, constraints));
+  }
+
+  Clock::pause();
+  Clock::settle();
+  Clock::advance(masterFlags.allocation_interval);
+  Clock::settle();
+
+  AWAIT_READY(offers);
+  EXPECT_EQ(offers->offers().size(), 1);
+
+  // TODO(asekretenko): After master starts exposing offer constraints via
+  // its endpoints (MESOS-10179), check the constraints in the endpoints.
+}
+
+
 } // namespace scheduler {
 } // namespace v1 {
 
