Skip to content

Commit

Permalink
add test
Browse files Browse the repository at this point in the history
Signed-off-by: Tim Vaillancourt <[email protected]>
  • Loading branch information
timvaillancourt committed Dec 12, 2024
1 parent 6e6f23a commit 4babcaf
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 6 deletions.
9 changes: 3 additions & 6 deletions go/vt/vtorc/logic/tablet_discovery.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
50 changes: 50 additions & 0 deletions go/vt/vtorc/logic/tablet_discovery_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ package logic
import (
"context"
"fmt"
"strings"
"sync/atomic"
"testing"
"time"
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 4babcaf

Please sign in to comment.