From 82eb2463798d45d28abea1cbc113ae8e84425979 Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Mon, 7 Apr 2025 17:29:22 -0400
Subject: [PATCH] Enhance tablets integration test to check stmt routing

Ensure that queries are routed according to learned tablet information.
---
 .../metadata/DefaultMetadataTabletMapIT.java  | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java
index bf772914f4b..6664909d0ec 100644
--- a/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java
+++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/metadata/DefaultMetadataTabletMapIT.java
@@ -20,6 +20,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -272,6 +273,13 @@ public void every_statement_should_deliver_tablet_info() {
               String.format(
                   "Statement %s on session %s did not trigger session tablets update",
                   stmtEntry.getKey(), sessionEntry.getKey()));
+          continue;
+        }
+        if (!checkIfRoutedProperly(session, stmt)) {
+          testErrors.add(
+              String.format(
+                  "Statement %s on session %s was routed to different nodes",
+                  stmtEntry.getKey(), sessionEntry.getKey()));
         }
       }
     }
@@ -341,6 +349,20 @@ private static boolean waitSessionLearnedTabletInfo(CqlSession session) {
     return isSessionLearnedTabletInfo(session);
   }
 
+  private static boolean checkIfRoutedProperly(CqlSession session, Statement stmt) {
+    // DefaultLoadBalancingPolicy suppose to prioritize nodes from replica list randomly shuffling
+    // them
+    // If routing goes wrong you will see more than REPLICATION_FACTOR unique first nodes in
+    // execution plan
+
+    int expectedNodesCount = stmt.isLWT() ? 1 : REPLICATION_FACTOR;
+    Set<Node> nodes = new HashSet<>();
+    for (int i = 0; i < REPLICATION_FACTOR * 3; i++) {
+      nodes.add(session.execute(stmt).getExecutionInfo().getCoordinator());
+    }
+    return nodes.size() <= expectedNodesCount;
+  }
+
   private static boolean isSessionLearnedTabletInfo(CqlSession session) {
     if (!session.getMetadata().getTabletMap().isPresent()) {
       return false;