vitess.io/vitess@v0.16.2/go/cmd/vtclient/vtclient_test.go (about)

     1  /*
     2  Copyright 2019 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 agreedto 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 main
    18  
    19  import (
    20  	"fmt"
    21  	"os"
    22  	"strings"
    23  	"testing"
    24  
    25  	"github.com/spf13/pflag"
    26  
    27  	"vitess.io/vitess/go/vt/vttest"
    28  
    29  	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
    30  	vttestpb "vitess.io/vitess/go/vt/proto/vttest"
    31  )
    32  
    33  func TestVtclient(t *testing.T) {
    34  	// Build the config for vttest.
    35  	var cfg vttest.Config
    36  	cfg.Topology = &vttestpb.VTTestTopology{
    37  		Keyspaces: []*vttestpb.Keyspace{
    38  			{
    39  				Name: "test_keyspace",
    40  				Shards: []*vttestpb.Shard{
    41  					{
    42  						Name: "0",
    43  					},
    44  				},
    45  			},
    46  		},
    47  	}
    48  	schema := `CREATE TABLE table1 (
    49  	  id BIGINT(20) UNSIGNED NOT NULL,
    50  	  i INT NOT NULL,
    51  	  PRIMARY KEY (id)
    52  	) ENGINE=InnoDB`
    53  	vschema := &vschemapb.Keyspace{
    54  		Sharded: true,
    55  		Vindexes: map[string]*vschemapb.Vindex{
    56  			"hash": {
    57  				Type: "hash",
    58  			},
    59  		},
    60  		Tables: map[string]*vschemapb.Table{
    61  			"table1": {
    62  				ColumnVindexes: []*vschemapb.ColumnVindex{
    63  					{
    64  						Column: "id",
    65  						Name:   "hash",
    66  					},
    67  				},
    68  			},
    69  		},
    70  	}
    71  	if err := cfg.InitSchemas("test_keyspace", schema, vschema); err != nil {
    72  		t.Fatalf("InitSchemas failed: %v", err)
    73  	}
    74  	defer os.RemoveAll(cfg.SchemaDir)
    75  	cluster := vttest.LocalCluster{
    76  		Config: cfg,
    77  	}
    78  	if err := cluster.Setup(); err != nil {
    79  		t.Fatalf("InitSchemas failed: %v", err)
    80  	}
    81  	defer cluster.TearDown() // nolint:errcheck
    82  
    83  	vtgateAddr := fmt.Sprintf("localhost:%v", cluster.Env.PortForProtocol("vtcombo", "grpc"))
    84  	queries := []struct {
    85  		args         []string
    86  		rowsAffected int64
    87  		errMsg       string
    88  	}{
    89  		{
    90  			args: []string{"SELECT * FROM table1"},
    91  		},
    92  		{
    93  			args: []string{"--target", "@primary", "--bind_variables", `[ 1, 100 ]`,
    94  				"INSERT INTO table1 (id, i) VALUES (:v1, :v2)"},
    95  			rowsAffected: 1,
    96  		},
    97  		{
    98  			args: []string{"--target", "@primary",
    99  				"UPDATE table1 SET i = (i + 1)"},
   100  			rowsAffected: 1,
   101  		},
   102  		{
   103  			args: []string{"--target", "@primary",
   104  				"SELECT * FROM table1"},
   105  			rowsAffected: 1,
   106  		},
   107  		{
   108  			args: []string{"--target", "@primary", "--bind_variables", `[ 1 ]`,
   109  				"DELETE FROM table1 WHERE id = :v1"},
   110  			rowsAffected: 1,
   111  		},
   112  		{
   113  			args: []string{"--target", "@primary",
   114  				"SELECT * FROM table1"},
   115  			rowsAffected: 0,
   116  		},
   117  		{
   118  			args:   []string{"SELECT * FROM nonexistent"},
   119  			errMsg: "table nonexistent not found",
   120  		},
   121  	}
   122  
   123  	// Change ErrorHandling from ExitOnError to panicking.
   124  	pflag.CommandLine.Init("vtclient_test.go", pflag.PanicOnError)
   125  	for _, q := range queries {
   126  		// Run main function directly and not as external process. To achieve this,
   127  		// overwrite os.Args which is used by pflag.Parse().
   128  		os.Args = []string{"vtclient_test.go", "--server", vtgateAddr}
   129  		os.Args = append(os.Args, q.args...)
   130  
   131  		results, err := run()
   132  		if q.errMsg != "" {
   133  			if got, want := err.Error(), q.errMsg; !strings.Contains(got, want) {
   134  				t.Fatalf("vtclient %v returned wrong error: got = %v, want contains = %v", os.Args[1:], got, want)
   135  			}
   136  			return
   137  		}
   138  
   139  		if err != nil {
   140  			t.Fatalf("vtclient %v failed: %v", os.Args[1:], err)
   141  		}
   142  		if got, want := results.rowsAffected, q.rowsAffected; got != want {
   143  			t.Fatalf("wrong rows affected for query: %v got = %v, want = %v", os.Args[1:], got, want)
   144  		}
   145  	}
   146  }