Adding tests to validate propagating local datacenter

Author: Michael Fero

cpp-driver/gtests/src/unit/tests/test_cluster.cpp

@@ -288,6 +288,37 @@ class ClusterUnitTest : public EventLoopTest {
     };
   };
 
+  class LocalDcClusterMetadataResolver : public ClusterMetadataResolver {
+  public:
+    LocalDcClusterMetadataResolver(const String& local_dc)
+        : desired_local_dc_(local_dc) {}
+
+  private:
+    virtual void internal_resolve(uv_loop_t* loop, const AddressVec& contact_points) {
+      resolved_contact_points_ = contact_points;
+      local_dc_ = desired_local_dc_;
+      callback_(this);
+    }
+
+    virtual void internal_cancel() {}
+
+  private:
+    String desired_local_dc_;
+  };
+
+  class LocalDcClusterMetadataResolverFactory : public ClusterMetadataResolverFactory {
+  public:
+    LocalDcClusterMetadataResolverFactory(const String& local_dc)
+        : local_dc_(local_dc) {}
+
+    virtual ClusterMetadataResolver::Ptr new_instance(const ClusterSettings& settings) const {
+      return ClusterMetadataResolver::Ptr(new LocalDcClusterMetadataResolver(local_dc_));
+    }
+
+  private:
+    String local_dc_;
+  };
+
   static void on_connection_connected(ClusterConnector* connector, Future* future) {
     if (connector->is_ok()) {
       future->set();
@@ -957,3 +988,25 @@ TEST_F(ClusterUnitTest, ReconnectionPolicy) {
   EXPECT_GE(policy->scheduled_delay_count(), 2u);
   EXPECT_EQ(3u, mock_cluster.connection_attempts(1)); // Includes initial connection attempt
 }
+
+TEST_F(ClusterUnitTest, LocalDcFromResolver) {
+  mockssandra::SimpleCluster cluster(simple(), 1);
+  ASSERT_EQ(cluster.start_all(), 0);
+
+  AddressVec contact_points;
+  contact_points.push_back(Address("127.0.0.1", 9042));
+
+  Future::Ptr connect_future(new Future());
+  ClusterConnector::Ptr connector(
+      new ClusterConnector(contact_points, PROTOCOL_VERSION,
+                           bind_callback(on_connection_reconnect, connect_future.get())));
+
+  ClusterSettings settings;
+  settings.cluster_metadata_resolver_factory = ClusterMetadataResolverFactory::Ptr(
+      new LocalDcClusterMetadataResolverFactory("this_local_dc"));
+  connector->with_settings(settings)->connect(event_loop());
+
+  ASSERT_TRUE(connect_future->wait_for(WAIT_FOR_TIME));
+  EXPECT_FALSE(connect_future->error());
+  ASSERT_EQ("this_local_dc", connect_future->cluster()->local_dc());
+}

cpp-driver/gtests/src/unit/tests/test_session.cpp

@@ -195,6 +195,37 @@ class SessionUnitTest : public EventLoopTest {
     uv_mutex_t mutex_;
     EventQueue events_;
   };
+
+  class LocalDcClusterMetadataResolver : public ClusterMetadataResolver {
+  public:
+    LocalDcClusterMetadataResolver(const String& local_dc)
+        : desired_local_dc_(local_dc) {}
+
+  private:
+    virtual void internal_resolve(uv_loop_t* loop, const AddressVec& contact_points) {
+      resolved_contact_points_ = contact_points;
+      local_dc_ = desired_local_dc_;
+      callback_(this);
+    }
+
+    virtual void internal_cancel() {}
+
+  private:
+    String desired_local_dc_;
+  };
+
+  class LocalDcClusterMetadataResolverFactory : public ClusterMetadataResolverFactory {
+  public:
+    LocalDcClusterMetadataResolverFactory(const String& local_dc)
+        : local_dc_(local_dc) {}
+
+    virtual ClusterMetadataResolver::Ptr new_instance(const ClusterSettings& settings) const {
+      return ClusterMetadataResolver::Ptr(new LocalDcClusterMetadataResolver(local_dc_));
+    }
+
+  private:
+    String local_dc_;
+  };
 };
 
 TEST_F(SessionUnitTest, ExecuteQueryNotConnected) {
@@ -610,3 +641,50 @@ TEST_F(SessionUnitTest, HostListenerNodeDown) {
 
   ASSERT_EQ(0u, listener->event_count());
 }
+
+TEST_F(SessionUnitTest, LocalDcUpdatedOnPolicy) {
+  mockssandra::SimpleCluster cluster(simple(), 3, 2);
+  ASSERT_EQ(cluster.start_all(), 0);
+
+  TestHostListener::Ptr listener(new TestHostListener());
+
+  Config config;
+  config.contact_points().push_back(Address("127.0.0.1", 9042));
+  config.set_use_randomized_contact_points(
+      false); // Ensure round robin order over DC for query execution
+  config.set_token_aware_routing(false);
+  config.set_load_balancing_policy(new DCAwarePolicy("dc1"));
+  config.set_cluster_metadata_resolver_factory(
+      ClusterMetadataResolverFactory::Ptr(new LocalDcClusterMetadataResolverFactory("dc2")));
+  config.set_host_listener(listener);
+
+  Session session;
+  connect(config, &session);
+
+  { // Initial nodes available from peers table (should skip DC1)
+    EXPECT_EQ(HostEventFuture::Event(HostEventFuture::ADD_NODE, Address("127.0.0.4", 9042)),
+              listener->wait_for_event(WAIT_FOR_TIME));
+    EXPECT_EQ(HostEventFuture::Event(HostEventFuture::START_NODE, Address("127.0.0.4", 9042)),
+              listener->wait_for_event(WAIT_FOR_TIME));
+    EXPECT_EQ(HostEventFuture::Event(HostEventFuture::ADD_NODE, Address("127.0.0.5", 9042)),
+              listener->wait_for_event(WAIT_FOR_TIME));
+    EXPECT_EQ(HostEventFuture::Event(HostEventFuture::START_NODE, Address("127.0.0.5", 9042)),
+              listener->wait_for_event(WAIT_FOR_TIME));
+  }
+
+  for (int i = 0; i < 20; ++i) { // Validate the request processors are using DC2 only
+    QueryRequest::Ptr request(new QueryRequest("blah", 0));
+
+    ResponseFuture::Ptr future = session.execute(request, NULL);
+    EXPECT_TRUE(future->wait_for(WAIT_FOR_TIME));
+    EXPECT_FALSE(future->error());
+
+    String expected_address = "127.0.0.";
+    expected_address.append((i % 2 == 0) ? "4" : "5");
+    EXPECT_EQ(expected_address, future->address().to_string());
+  }
+
+  close(&session);
+
+  ASSERT_EQ(0u, listener->event_count());
+}