vitess.io/vitess@v0.16.2/go/vt/wrangler/testlib/vtctl_pipe.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 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 testlib
    18  
    19  import (
    20  	"bytes"
    21  	"context"
    22  	"fmt"
    23  	"io"
    24  	"net"
    25  	"sync"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/spf13/pflag"
    30  	"github.com/stretchr/testify/require"
    31  	"google.golang.org/grpc"
    32  
    33  	"vitess.io/vitess/go/vt/logutil"
    34  	"vitess.io/vitess/go/vt/servenv"
    35  	"vitess.io/vitess/go/vt/topo"
    36  	"vitess.io/vitess/go/vt/vtctl/grpcvtctlserver"
    37  	"vitess.io/vitess/go/vt/vtctl/vtctlclient"
    38  
    39  	// we need to import the grpcvtctlclient library so the gRPC
    40  	// vtctl client is registered and can be used.
    41  	_ "vitess.io/vitess/go/vt/vtctl/grpcvtctlclient"
    42  )
    43  
    44  var servenvInitialized sync.Once
    45  
    46  // VtctlPipe is a vtctl server based on a topo server, and a client that
    47  // is connected to it via gRPC.
    48  type VtctlPipe struct {
    49  	listener net.Listener
    50  	client   vtctlclient.VtctlClient
    51  	t        *testing.T
    52  }
    53  
    54  // NewVtctlPipe creates a new VtctlPipe based on the given topo server.
    55  func NewVtctlPipe(t *testing.T, ts *topo.Server) *VtctlPipe {
    56  	// Register all vtctl commands
    57  	servenvInitialized.Do(func() {
    58  		// make sure we use the right protocol
    59  		fs := pflag.NewFlagSet("", pflag.ContinueOnError)
    60  		vtctlclient.RegisterFlags(fs)
    61  
    62  		err := fs.Parse([]string{
    63  			"--vtctl_client_protocol",
    64  			"grpc",
    65  		})
    66  		require.NoError(t, err, "failed to set `--vtctl_client_protocol=%s`", "grpc")
    67  
    68  		servenv.FireRunHooks()
    69  	})
    70  
    71  	// Listen on a random port
    72  	listener, err := net.Listen("tcp", "127.0.0.1:0")
    73  	if err != nil {
    74  		t.Fatalf("Cannot listen: %v", err)
    75  	}
    76  
    77  	// Create a gRPC server and listen on the port
    78  	server := grpc.NewServer()
    79  	grpcvtctlserver.StartServer(server, ts)
    80  	go server.Serve(listener)
    81  
    82  	// Create a VtctlClient gRPC client to talk to the fake server
    83  	client, err := vtctlclient.New(listener.Addr().String())
    84  	if err != nil {
    85  		t.Fatalf("Cannot create client: %v", err)
    86  	}
    87  
    88  	return &VtctlPipe{
    89  		listener: listener,
    90  		client:   client,
    91  		t:        t,
    92  	}
    93  }
    94  
    95  // Close will stop listening and free up all resources.
    96  func (vp *VtctlPipe) Close() {
    97  	vp.client.Close()
    98  	vp.listener.Close()
    99  }
   100  
   101  // Run executes the provided command remotely, logs the output in the
   102  // test logs, and returns the command error.
   103  func (vp *VtctlPipe) Run(args []string) error {
   104  	return vp.run(args, func(line string) {
   105  		vp.t.Log(line)
   106  	})
   107  }
   108  
   109  // RunAndOutput is similar to Run, but returns the output as a multi-line string
   110  // instead of logging it.
   111  func (vp *VtctlPipe) RunAndOutput(args []string) (string, error) {
   112  	var output bytes.Buffer
   113  	err := vp.run(args, func(line string) {
   114  		output.WriteString(line)
   115  	})
   116  	return output.String(), err
   117  }
   118  
   119  func (vp *VtctlPipe) run(args []string, outputFunc func(string)) error {
   120  	actionTimeout := 30 * time.Second
   121  	ctx := context.Background()
   122  
   123  	stream, err := vp.client.ExecuteVtctlCommand(ctx, args, actionTimeout)
   124  	if err != nil {
   125  		return fmt.Errorf("VtctlPipe.Run() failed: %v", err)
   126  	}
   127  	for {
   128  		le, err := stream.Recv()
   129  		switch err {
   130  		case nil:
   131  			outputFunc(logutil.EventString(le))
   132  		case io.EOF:
   133  			return nil
   134  		default:
   135  			return err
   136  		}
   137  	}
   138  }
   139  
   140  // RunAndStreamOutput returns the output of the vtctl command as a channel.
   141  // When the channcel is closed, the command did finish.
   142  func (vp *VtctlPipe) RunAndStreamOutput(args []string) (logutil.EventStream, error) {
   143  	actionTimeout := 30 * time.Second
   144  	ctx := context.Background()
   145  
   146  	return vp.client.ExecuteVtctlCommand(ctx, args, actionTimeout)
   147  }