vitess.io/vitess@v0.16.2/go/vt/vtctl/grpcvtctldclient/client_test.go (about)

     1  /*
     2  Copyright 2021 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package grpcvtctldclient
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/stretchr/testify/require"
    25  
    26  	"vitess.io/vitess/go/test/utils"
    27  	"vitess.io/vitess/go/vt/topo"
    28  	"vitess.io/vitess/go/vt/topo/memorytopo"
    29  	"vitess.io/vitess/go/vt/vtctl/grpcvtctldserver"
    30  	"vitess.io/vitess/go/vt/vtctl/grpcvtctldserver/testutil"
    31  	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
    32  
    33  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    34  	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
    35  	vtctlservicepb "vitess.io/vitess/go/vt/proto/vtctlservice"
    36  )
    37  
    38  func TestFindAllShardsInKeyspace(t *testing.T) {
    39  	ctx := context.Background()
    40  	ts := memorytopo.NewServer("cell1")
    41  	vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer {
    42  		return grpcvtctldserver.NewVtctldServer(ts)
    43  	})
    44  
    45  	testutil.WithTestServer(t, vtctld, func(t *testing.T, client vtctldclient.VtctldClient) {
    46  		ks := &vtctldatapb.Keyspace{
    47  			Name:     "testkeyspace",
    48  			Keyspace: &topodatapb.Keyspace{},
    49  		}
    50  		testutil.AddKeyspace(ctx, t, ts, ks)
    51  
    52  		si1, err := ts.GetOrCreateShard(ctx, ks.Name, "-80")
    53  		require.NoError(t, err)
    54  		si2, err := ts.GetOrCreateShard(ctx, ks.Name, "80-")
    55  		require.NoError(t, err)
    56  
    57  		resp, err := client.FindAllShardsInKeyspace(ctx, &vtctldatapb.FindAllShardsInKeyspaceRequest{Keyspace: ks.Name})
    58  		assert.NoError(t, err)
    59  		assert.NotNil(t, resp)
    60  
    61  		expected := map[string]*vtctldatapb.Shard{
    62  			"-80": {
    63  				Keyspace: ks.Name,
    64  				Name:     "-80",
    65  				Shard:    si1.Shard,
    66  			},
    67  			"80-": {
    68  				Keyspace: ks.Name,
    69  				Name:     "80-",
    70  				Shard:    si2.Shard,
    71  			},
    72  		}
    73  
    74  		utils.MustMatch(t, expected, resp.Shards)
    75  
    76  		client.Close()
    77  		_, err = client.FindAllShardsInKeyspace(ctx, &vtctldatapb.FindAllShardsInKeyspaceRequest{Keyspace: ks.Name})
    78  		assert.Error(t, err)
    79  	})
    80  }
    81  
    82  func TestGetKeyspace(t *testing.T) {
    83  	ctx := context.Background()
    84  
    85  	ts := memorytopo.NewServer("cell1")
    86  	vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer {
    87  		return grpcvtctldserver.NewVtctldServer(ts)
    88  	})
    89  
    90  	testutil.WithTestServer(t, vtctld, func(t *testing.T, client vtctldclient.VtctldClient) {
    91  		expected := &vtctldatapb.GetKeyspaceResponse{
    92  			Keyspace: &vtctldatapb.Keyspace{
    93  				Name:     "testkeyspace",
    94  				Keyspace: &topodatapb.Keyspace{},
    95  			},
    96  		}
    97  		testutil.AddKeyspace(ctx, t, ts, expected.Keyspace)
    98  
    99  		resp, err := client.GetKeyspace(ctx, &vtctldatapb.GetKeyspaceRequest{Keyspace: expected.Keyspace.Name})
   100  		assert.NoError(t, err)
   101  		utils.MustMatch(t, expected, resp)
   102  
   103  		client.Close()
   104  		_, err = client.GetKeyspace(ctx, &vtctldatapb.GetKeyspaceRequest{})
   105  		assert.Error(t, err)
   106  	})
   107  }
   108  
   109  func TestGetKeyspaces(t *testing.T) {
   110  	ctx := context.Background()
   111  
   112  	ts := memorytopo.NewServer("cell1")
   113  	vtctld := testutil.NewVtctldServerWithTabletManagerClient(t, ts, nil, func(ts *topo.Server) vtctlservicepb.VtctldServer {
   114  		return grpcvtctldserver.NewVtctldServer(ts)
   115  	})
   116  
   117  	testutil.WithTestServer(t, vtctld, func(t *testing.T, client vtctldclient.VtctldClient) {
   118  		resp, err := client.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
   119  		assert.NoError(t, err)
   120  		assert.Empty(t, resp.Keyspaces)
   121  
   122  		expected := &vtctldatapb.Keyspace{
   123  			Name:     "testkeyspace",
   124  			Keyspace: &topodatapb.Keyspace{},
   125  		}
   126  		testutil.AddKeyspace(ctx, t, ts, expected)
   127  
   128  		resp, err = client.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
   129  		assert.NoError(t, err)
   130  		utils.MustMatch(t, []*vtctldatapb.Keyspace{expected}, resp.Keyspaces)
   131  
   132  		client.Close()
   133  		_, err = client.GetKeyspaces(ctx, &vtctldatapb.GetKeyspacesRequest{})
   134  		assert.Error(t, err)
   135  	})
   136  }