[5374] Checkpoint: HA tests

1 files changed, 27 insertions, 1 deletions

diff --git a/src/bin/dhcp4/tests/dhcp4_client.cc b/src/bin/dhcp4/tests/dhcp4_client.cc
index c117ba4320..8475989105 100644
--- a/src/bin/dhcp4/tests/dhcp4_client.cc
+++ b/src/bin/dhcp4/tests/dhcp4_client.cc
@@ -1,4 +1,4 @@
-// Copyright (C) 2014-2017 Internet Systems Consortium, Inc. ("ISC")
+// Copyright (C) 2014-2018 Internet Systems Consortium, Inc. ("ISC")
 //
 // This Source Code Form is subject to the terms of the Mozilla Public
 // License, v. 2.0. If a copy of the MPL was not distributed with this
@@ -239,6 +239,14 @@ Dhcp4Client::appendExtraOptions() {
 }
 
 void
+Dhcp4Client::appendClasses() {
+    for (ClientClasses::const_iterator cclass = classes_.cbegin();
+         cclass != classes_.cend(); ++cclass) {
+        context_.query_->addClass(*cclass);
+    }
+}
+
+void
 Dhcp4Client::doDiscover(const boost::shared_ptr<IOAddress>& requested_addr) {
     context_.query_ = createMsg(DHCPDISCOVER);
     // Request options if any.
@@ -255,6 +263,7 @@ Dhcp4Client::doDiscover(const boost::shared_ptr<IOAddress>& requested_addr) {
         context_.query_->setCiaddr(ciaddr_.get());
     }
     appendExtraOptions();
+    appendClasses();
 
     // Send the message to the server.
     sendMsg(context_.query_);
@@ -277,6 +286,8 @@ Dhcp4Client::doInform(const bool set_ciaddr) {
     appendPRL();
     // Any other options to be sent by a client.
     appendExtraOptions();
+    // Add classes.
+    appendClasses();
     // The client sending a DHCPINFORM message has an IP address obtained
     // by some other means, e.g. static configuration. The lease which we
     // are using here is most likely set by the createLease method.
@@ -396,6 +407,8 @@ Dhcp4Client::doRequest() {
     appendClientId();
     // Any other options to be sent by a client.
     appendExtraOptions();
+    // Add classes.
+    appendClasses();
     // Send the message to the server.
     sendMsg(context_.query_);
     // Expect response.
@@ -518,6 +531,12 @@ Dhcp4Client::sendMsg(const Pkt4Ptr& msg) {
     msg_copy->setRemoteAddr(msg->getLocalAddr());
     msg_copy->setLocalAddr(dest_addr_);
     msg_copy->setIface(iface_name_);
+    // Copy classes
+    const ClientClasses& classes = msg->getClasses();
+    for (ClientClasses::const_iterator cclass = classes.cbegin();
+         cclass != classes.cend(); ++cclass) {
+        msg_copy->addClass(*cclass);
+    }
     srv_->fakeReceive(msg_copy);
     srv_->run();
 }
@@ -536,6 +555,13 @@ Dhcp4Client::addExtraOption(const OptionPtr& opt) {
     extra_options_.insert(std::make_pair(opt->getType(), opt));
 }
 
+void
+Dhcp4Client::addClass(const ClientClass& client_class) {
+    if (!classes_.contains(client_class)) {
+        classes_.insert(client_class);
+    }
+}
+
 } // end of namespace isc::dhcp::test
 } // end of namespace isc::dhcp
 } // end of namespace isc