Make ringDescriber GetHosts() testable and add test for it

scylladb · Dec 6, 2024 · 0d743ce · 0d743ce
1 parent 24e876c
commit 0d743ce
Show file tree

Hide file tree

Showing 3 changed files with 387 additions and 94 deletions.
diff --git a/host_source_test.go b/host_source_test.go
@@ -364,81 +364,3 @@ func TestErrorBroadcaster_StopWithoutBroadcast(t *testing.T) {
 		t.Errorf(loadedVal.(error).Error())
 	}
 }
-
-func TestGetClusterPeerInfoZeroToken(t *testing.T) {
-	host_id1, _ := ParseUUID("b2035fd9-e0ca-4857-8c45-e63c00fb7c43")
-	host_id2, _ := ParseUUID("4b21ee4c-acea-4267-8e20-aaed5361a0dd")
-	host_id3, _ := ParseUUID("dfef4a22-b8d8-47e9-aee5-8c19d4b7a9e3")
-
-	schema_version1, _ := ParseUUID("af810386-a694-11ef-81fa-3aea73156247")
-
-	peersRows := []map[string]interface{}{
-		{
-			"data_center":     "datacenter1",
-			"host_id":         host_id1,
-			"peer":            "127.0.0.3",
-			"preferred_ip":    "127.0.0.3",
-			"rack":            "rack1",
-			"release_version": "3.0.8",
-			"rpc_address":     "127.0.0.3",
-			"schema_version":  schema_version1,
-			"tokens":          []string{"-1296227678594315580994457470329811265"},
-		},
-		{
-			"data_center":     "datacenter1",
-			"host_id":         host_id2,
-			"peer":            "127.0.0.2",
-			"preferred_ip":    "127.0.0.2",
-			"rack":            "rack1",
-			"release_version": "3.0.8",
-			"rpc_address":     "127.0.0.2",
-			"schema_version":  schema_version1,
-			"tokens":          []string{"-1129762924682054333"},
-		},
-		{
-			"data_center":     "datacenter2",
-			"host_id":         host_id3,
-			"peer":            "127.0.0.5",
-			"preferred_ip":    "127.0.0.5",
-			"rack":            "rack1",
-			"release_version": "3.0.8",
-			"rpc_address":     "127.0.0.5",
-			"schema_version":  schema_version1,
-			"tokens":          []string{},
-		},
-	}
-
-	translateAddressPort := func(addr net.IP, port int) (net.IP, int) {
-		return addr, port
-	}
-
-	var logger StdLogger
-	t.Run("OmitOneZeroTokenNode", func(t *testing.T) {
-		peers, err := getPeersFromQuerySystemPeers(
-			peersRows,
-			9042,
-			translateAddressPort,
-			logger,
-		)
-
-		if err != nil {
-			t.Fatalf("unable to get peers: %v", err)
-		}
-		assertEqual(t, "peers length", 2, len(peers))
-	})
-
-	t.Run("NoZeroTokenNodes", func(t *testing.T) {
-		peersRows[2]["tokens"] = []string{"-1129762924682054333"}
-		peers, err := getPeersFromQuerySystemPeers(
-			peersRows,
-			9042,
-			translateAddressPort,
-			logger,
-		)
-
-		if err != nil {
-			t.Fatalf("unable to get peers: %v", err)
-		}
-		assertEqual(t, "peers length", 3, len(peers))
-	})
-}
diff --git a/ring_describer.go b/ring_describer.go
@@ -26,13 +26,8 @@ func (r *ringDescriber) setControlConn(c controlConnection) {
 }
 
 // Ask the control node for the local host info
-func (r *ringDescriber) getLocalHostInfo() (*HostInfo, error) {
-	if r.control == nil {
-		return nil, errNoControl
-	}
-
-	ch := r.control.getConn()
-	iter := querySystemLocal(context.TODO(), ch.conn)
+func (r *ringDescriber) getLocalHostInfo(conn ConnInterface) (*HostInfo, error) {
+	iter := querySystemLocal(context.TODO(), conn)
 
 	if iter == nil {
 		return nil, errNoControl
@@ -46,13 +41,8 @@ func (r *ringDescriber) getLocalHostInfo() (*HostInfo, error) {
 }
 
 // Ask the control node for host info on all it's known peers
-func (r *ringDescriber) getClusterPeerInfo(localHost *HostInfo) ([]*HostInfo, error) {
-	if r.control == nil {
-		return nil, errNoControl
-	}
-
-	ch := r.control.getConn()
-	iter := querySystemPeers(context.TODO(), ch.conn)
+func (r *ringDescriber) getClusterPeerInfo(localHost *HostInfo, c ConnInterface) ([]*HostInfo, error) {
+	iter := querySystemPeers(context.TODO(), c)
 
 	if iter == nil {
 		return nil, errNoControl
@@ -107,12 +97,17 @@ func (r *ringDescriber) GetHosts() ([]*HostInfo, string, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	localHost, err := r.getLocalHostInfo()
+	if r.control == nil {
+		return r.prevHosts, r.prevPartitioner, errNoControl
+	}
+
+	ch := r.control.getConn()
+	localHost, err := r.getLocalHostInfo(ch.conn)
 	if err != nil {
 		return r.prevHosts, r.prevPartitioner, err
 	}
 
-	peerHosts, err := r.getClusterPeerInfo(localHost)
+	peerHosts, err := r.getClusterPeerInfo(localHost, ch.conn)
 	if err != nil {
 		return r.prevHosts, r.prevPartitioner, err
 	}