From 4babcaf2e3b49ff51b635dd95df5d1a9f0016e76 Mon Sep 17 00:00:00 2001
From: Tim Vaillancourt <tim@timvaillancourt.com>
Date: Thu, 12 Dec 2024 23:31:30 +0100
Subject: [PATCH] add test

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
---
 go/vt/vtorc/logic/tablet_discovery.go      |  9 ++--
 go/vt/vtorc/logic/tablet_discovery_test.go | 50 ++++++++++++++++++++++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/go/vt/vtorc/logic/tablet_discovery.go b/go/vt/vtorc/logic/tablet_discovery.go
index e5aed1fb41e..c56836f089d 100644
--- a/go/vt/vtorc/logic/tablet_discovery.go
+++ b/go/vt/vtorc/logic/tablet_discovery.go
@@ -34,7 +34,6 @@ import (
 	"vitess.io/vitess/go/vt/external/golib/sqlutils"
 	"vitess.io/vitess/go/vt/log"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
-	"vitess.io/vitess/go/vt/servenv"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtorc/config"
@@ -49,16 +48,12 @@ var (
 	clustersToWatch   []string
 	shutdownWaitTime  = 30 * time.Second
 	shardsLockCounter int32
-	shardsToWatch     map[string]bool
+	shardsToWatch     = make(map[string]bool, 0)
 
 	// ErrNoPrimaryTablet is a fixed error message.
 	ErrNoPrimaryTablet = errors.New("no primary tablet found")
 )
 
-func init() {
-	servenv.OnInit(parseClustersToWatch)
-}
-
 // parseClustersToWatch parses the --clusters_to_watch flag-value
 // into a map of keyspace/shards. This is called once at init
 // time because the list never changes.
@@ -131,6 +126,8 @@ func OpenTabletDiscovery() <-chan time.Time {
 	if _, err := db.ExecVTOrc("DELETE FROM vitess_tablet"); err != nil {
 		log.Error(err)
 	}
+	// Parse --clusters_to_watch into a filter.
+	parseClustersToWatch()
 	// We refresh all information from the topo once before we start the ticks to do
 	// it on a timer.
 	ctx, cancel := context.WithTimeout(context.Background(), topo.RemoteOperationTimeout)
diff --git a/go/vt/vtorc/logic/tablet_discovery_test.go b/go/vt/vtorc/logic/tablet_discovery_test.go
index f6a7af38382..cbb39bc8c71 100644
--- a/go/vt/vtorc/logic/tablet_discovery_test.go
+++ b/go/vt/vtorc/logic/tablet_discovery_test.go
@@ -19,6 +19,7 @@ package logic
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -101,6 +102,55 @@ var (
 	}
 )
 
+func TestParseClustersToWatch(t *testing.T) {
+	oldClustersToWatch := clustersToWatch
+	oldTs := ts
+	defer func() {
+		clustersToWatch = oldClustersToWatch
+		shardsToWatch = nil
+		ts = oldTs
+	}()
+
+	// Create a memory topo-server and create the keyspace and shard records
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ts = memorytopo.NewServer(ctx, cell1)
+	_, err := ts.GetOrCreateShard(context.Background(), keyspace, shard)
+	require.NoError(t, err)
+
+	testCases := []struct {
+		in       []string
+		expected map[string]bool
+	}{
+		{
+			in:       []string{"test/"},
+			expected: map[string]bool{"test/": true},
+		},
+		{
+			in:       []string{"test/-"},
+			expected: map[string]bool{"test/-": true},
+		},
+		{
+			in: []string{keyspace},
+			expected: map[string]bool{
+				topoproto.KeyspaceShardString(keyspace, shard): true,
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(strings.Join(testCase.in, ","), func(t *testing.T) {
+			defer func() {
+				shardsToWatch = make(map[string]bool, 0)
+			}()
+			clustersToWatch = testCase.in
+			parseClustersToWatch()
+			require.Equal(t, testCase.expected, shardsToWatch)
+		})
+	}
+}
+
 func TestRefreshTabletsInKeyspaceShard(t *testing.T) {
 	// Store the old flags and restore on test completion
 	oldTs := ts