test: make it easier to run tests without a main keyspace (#17501)

Signed-off-by: Andres Taylor <[email protected]>

go/test/vschemawrapper/vschema_wrapper.go

@@ -18,7 +18,9 @@ package vschemawrapper
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"sort"
 	"strings"
 
 	"vitess.io/vitess/go/mysql/collations"
@@ -275,17 +277,6 @@ func (vw *VSchemaWrapper) FindTableOrVindex(tab sqlparser.TableName) (*vindexes.
 	return vw.Vcursor.FindTableOrVindex(tab)
 }
 
-func (vw *VSchemaWrapper) getfirstKeyspace() (ks *vindexes.Keyspace) {
-	var f string
-	for name, schema := range vw.V.Keyspaces {
-		if f == "" || f > name {
-			f = name
-			ks = schema.Keyspace
-		}
-	}
-	return
-}
-
 func (vw *VSchemaWrapper) getActualKeyspace() string {
 	if vw.Keyspace == nil {
 		return ""
@@ -301,15 +292,32 @@ func (vw *VSchemaWrapper) getActualKeyspace() string {
 }
 
 func (vw *VSchemaWrapper) SelectedKeyspace() (*vindexes.Keyspace, error) {
-	return vw.V.Keyspaces["main"].Keyspace, nil
+	return vw.AnyKeyspace()
 }
 
 func (vw *VSchemaWrapper) AnyKeyspace() (*vindexes.Keyspace, error) {
-	return vw.SelectedKeyspace()
+	ks, found := vw.V.Keyspaces["main"]
+	if found {
+		return ks.Keyspace, nil
+	}
+
+	size := len(vw.V.Keyspaces)
+	if size == 0 {
+		return nil, errors.New("no keyspace found in vschema")
+	}
+
+	// Find the first keyspace in the map alphabetically to get deterministic results
+	keys := make([]string, size)
+	for key := range vw.V.Keyspaces {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+
+	return vw.V.Keyspaces[keys[0]].Keyspace, nil
 }
 
 func (vw *VSchemaWrapper) FirstSortedKeyspace() (*vindexes.Keyspace, error) {
-	return vw.V.Keyspaces["main"].Keyspace, nil
+	return vw.AnyKeyspace()
 }
 
 func (vw *VSchemaWrapper) TargetString() string {
@@ -344,12 +352,12 @@ func (vw *VSchemaWrapper) IsViewsEnabled() bool {
 
 // FindMirrorRule finds the mirror rule for the requested keyspace, table
 // name, and the tablet type in the VSchema.
-func (vs *VSchemaWrapper) FindMirrorRule(tab sqlparser.TableName) (*vindexes.MirrorRule, error) {
+func (vw *VSchemaWrapper) FindMirrorRule(tab sqlparser.TableName) (*vindexes.MirrorRule, error) {
 	destKeyspace, destTabletType, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return nil, err
 	}
-	mirrorRule, err := vs.V.FindMirrorRule(destKeyspace, tab.Name.String(), destTabletType)
+	mirrorRule, err := vw.V.FindMirrorRule(destKeyspace, tab.Name.String(), destTabletType)
 	if err != nil {
 		return nil, err
 	}