vitess.io/vitess@v0.16.2/go/test/endtoend/vtcombo/recreate/recreate_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 recreate
    18  
    19  import (
    20  	"context"
    21  	"flag"
    22  	"fmt"
    23  	"os"
    24  	"os/exec"
    25  	"strings"
    26  	"testing"
    27  
    28  	"github.com/stretchr/testify/assert"
    29  	"github.com/stretchr/testify/require"
    30  
    31  	"vitess.io/vitess/go/vt/log"
    32  	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
    33  	"vitess.io/vitess/go/vt/vttest"
    34  
    35  	vttestpb "vitess.io/vitess/go/vt/proto/vttest"
    36  )
    37  
    38  var (
    39  	localCluster *vttest.LocalCluster
    40  	grpcAddress  string
    41  	vtctldAddr   string
    42  	ks1          = "test_keyspace"
    43  	redirected   = "redirected"
    44  )
    45  
    46  func TestMain(m *testing.M) {
    47  	flag.Parse()
    48  
    49  	exitcode, err := func() (int, error) {
    50  
    51  		topology := new(vttestpb.VTTestTopology)
    52  		topology.Keyspaces = []*vttestpb.Keyspace{
    53  			{
    54  				Name: ks1,
    55  				Shards: []*vttestpb.Shard{
    56  					{Name: "-80"},
    57  					{Name: "80-"},
    58  				},
    59  				RdonlyCount:  1,
    60  				ReplicaCount: 2,
    61  			},
    62  			{
    63  				Name:       redirected,
    64  				ServedFrom: ks1,
    65  			},
    66  		}
    67  
    68  		var cfg vttest.Config
    69  		cfg.Topology = topology
    70  		cfg.SchemaDir = os.Getenv("VTROOT") + "/test/vttest_schema"
    71  		cfg.DefaultSchemaDir = os.Getenv("VTROOT") + "/test/vttest_schema/default"
    72  		cfg.PersistentMode = true
    73  
    74  		localCluster = &vttest.LocalCluster{
    75  			Config: cfg,
    76  		}
    77  
    78  		err := localCluster.Setup()
    79  		defer localCluster.TearDown()
    80  		if err != nil {
    81  			return 1, err
    82  		}
    83  
    84  		grpcAddress = fmt.Sprintf("localhost:%d", localCluster.Env.PortForProtocol("vtcombo", "grpc"))
    85  		vtctldAddr = fmt.Sprintf("localhost:%d", localCluster.Env.PortForProtocol("vtcombo", "port"))
    86  
    87  		return m.Run(), nil
    88  	}()
    89  	if err != nil {
    90  		log.Errorf("top level error: %v\n", err)
    91  		os.Exit(1)
    92  	} else {
    93  		os.Exit(exitcode)
    94  	}
    95  }
    96  
    97  func TestDropAndRecreateWithSameShards(t *testing.T) {
    98  	ctx := context.Background()
    99  	conn, err := vtgateconn.Dial(ctx, grpcAddress)
   100  	require.Nil(t, err)
   101  	defer conn.Close()
   102  
   103  	cur := conn.Session(ks1+"@primary", nil)
   104  
   105  	_, err = cur.Execute(ctx, "DROP DATABASE "+ks1, nil)
   106  	require.Nil(t, err)
   107  
   108  	_, err = cur.Execute(ctx, "CREATE DATABASE "+ks1, nil)
   109  	require.Nil(t, err)
   110  
   111  	assertTabletsPresent(t)
   112  }
   113  
   114  func assertTabletsPresent(t *testing.T) {
   115  	tmpCmd := exec.Command("vtctlclient", "--vtctl_client_protocol", "grpc", "--server", grpcAddress, "--stderrthreshold", "0", "ListAllTablets", "--", "test")
   116  
   117  	log.Infof("Running vtctlclient with command: %v", tmpCmd.Args)
   118  
   119  	output, err := tmpCmd.CombinedOutput()
   120  	require.Nil(t, err)
   121  
   122  	numPrimary, numReplica, numRdonly, numDash80, num80Dash := 0, 0, 0, 0, 0
   123  	lines := strings.Split(string(output), "\n")
   124  	for _, line := range lines {
   125  		if !strings.HasPrefix(line, "test-") {
   126  			continue
   127  		}
   128  		parts := strings.Split(line, " ")
   129  		assert.Equal(t, "test_keyspace", parts[1])
   130  
   131  		switch parts[3] {
   132  		case "primary":
   133  			numPrimary++
   134  		case "replica":
   135  			numReplica++
   136  		case "rdonly":
   137  			numRdonly++
   138  		default:
   139  			t.Logf("invalid tablet type %s", parts[3])
   140  		}
   141  
   142  		switch parts[2] {
   143  		case "-80":
   144  			numDash80++
   145  		case "80-":
   146  			num80Dash++
   147  		default:
   148  			t.Logf("invalid shard %s", parts[2])
   149  		}
   150  
   151  	}
   152  
   153  	assert.Equal(t, 2, numPrimary)
   154  	assert.Equal(t, 2, numReplica)
   155  	assert.Equal(t, 2, numRdonly)
   156  	assert.Equal(t, 3, numDash80)
   157  	assert.Equal(t, 3, num80Dash)
   158  }