func commandSucceeds(t *testing.T, client vtworkerclient.Client) {
logs, errFunc, err := client.ExecuteVtworkerCommand(context.Background(), []string{"Ping", "pong"})
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
count := 0
for e := range logs {
expected := "Ping command was called with message: 'pong'.\n"
if logutil.EventString(e) != expected {
t.Errorf("Got unexpected log line '%v' expected '%v'", e.String(), expected)
}
count++
}
if count != 1 {
t.Errorf("Didn't get expected log line only, got %v lines", count)
}
if err := errFunc(); err != nil {
t.Fatalf("Remote error: %v", err)
}
logs, errFunc, err = client.ExecuteVtworkerCommand(context.Background(), []string{"Reset"})
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
for range logs {
}
if err := errFunc(); err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
}
func commandPanics(t *testing.T, client vtworkerclient.Client) {
stream, err := client.ExecuteVtworkerCommand(context.Background(), []string{"Panic"})
// The expected error could already be seen now or after the output channel is closed.
// To avoid checking for the same error twice, we don't check it here yet.
if err == nil {
// Don't check for errors until the output channel is closed.
// No output expected in this case.
_, err = stream.Recv()
}
expected := "uncaught vtworker panic: Panic command was called. This should be caught by the vtworker framework and logged as an error."
if err == nil || !strings.Contains(err.Error(), expected) {
t.Fatalf("Unexpected remote error, got: '%v' was expecting to find '%v'", err, expected)
}
}
func runVtworkerCommand(client vtworkerclient.Client, args []string) error {
stream, err := client.ExecuteVtworkerCommand(context.Background(), args)
if err != nil {
return fmt.Errorf("cannot execute remote command: %v", err)
}
for {
_, err := stream.Recv()
switch err {
case nil:
// Consume next response.
case io.EOF:
return nil
default:
return vterrors.WithPrefix("unexpected error when reading the stream: ", err)
}
}
}
func commandErrors(t *testing.T, client vtworkerclient.Client) {
logs, errFunc, err := client.ExecuteVtworkerCommand(context.Background(), []string{"NonexistingCommand"})
// The expected error could already be seen now or after the output channel is closed.
// To avoid checking for the same error twice, we don't check it here yet.
if err == nil {
// Don't check for errors until the output channel is closed.
// We expect the usage to be sent as output. However, we have to consider it
// optional and do not test for it because not all RPC implementations send
// the output after an error.
for {
if _, ok := <-logs; !ok {
break
}
}
err = errFunc()
}
expected := "unknown command: NonexistingCommand"
if err == nil || !strings.Contains(err.Error(), expected) {
t.Fatalf("Unexpected remote error, got: '%v' was expecting to find '%v'", err, expected)
}
}
func commandSucceeds(t *testing.T, client vtworkerclient.Client) {
stream, err := client.ExecuteVtworkerCommand(context.Background(), []string{"Ping", "pong"})
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
got, err := stream.Recv()
if err != nil {
t.Fatalf("failed to get first line: %v", err)
}
expected := "Ping command was called with message: 'pong'.\n"
if logutil.EventString(got) != expected {
t.Errorf("Got unexpected log line '%v' expected '%v'", got.String(), expected)
}
got, err = stream.Recv()
if err != io.EOF {
t.Fatalf("Didn't get EOF as expected: %v", err)
}
stream, err = client.ExecuteVtworkerCommand(context.Background(), []string{"Reset"})
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
for {
_, err := stream.Recv()
switch err {
case nil:
// next please!
case io.EOF:
// done with test
return
default:
// unexpected error
t.Fatalf("Cannot execute remote command: %v", err)
}
}
}
func commandSucceeds(t *testing.T, client vtworkerclient.Client) {
stream, err := client.ExecuteVtworkerCommand(context.Background(), []string{"Ping", "pong"})
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
got, err := stream.Recv()
if err != nil {
t.Fatalf("failed to get first line: %v", err)
}
expected := "Ping command was called with message: 'pong'.\n"
if logutil.EventString(got) != expected {
t.Errorf("Got unexpected log line '%v' expected '%v'", got.String(), expected)
}
got, err = stream.Recv()
if err != io.EOF {
t.Fatalf("Didn't get EOF as expected: %v", err)
}
// Reset vtworker for the next test function.
if err := runVtworkerCommand(client, []string{"Reset"}); err != nil {
t.Fatal(err)
}
}
// commandErrorsBecauseBusy tests that concurrent commands are rejected with
// TRANSIENT_ERROR while a command is already running.
// It also tests the correct propagation of the CANCELED error code.
func commandErrorsBecauseBusy(t *testing.T, client vtworkerclient.Client, serverSideCancelation bool) {
// Run the vtworker "Block" command which blocks until we cancel the context.
var wg sync.WaitGroup
ctx, cancel := context.WithCancel(context.Background())
// blockCommandStarted will be closed after we're sure that vtworker is
// running the "Block" command.
blockCommandStarted := make(chan struct{})
var errorCodeCheck error
wg.Add(1)
go func() {
stream, err := client.ExecuteVtworkerCommand(ctx, []string{"Block"})
if err != nil {
t.Fatalf("Block command should not have failed: %v", err)
}
firstLineReceived := false
for {
if _, err := stream.Recv(); err != nil {
// We see CANCELED from the RPC client (client side cancelation) or
// from vtworker itself (server side cancelation).
if vterrors.RecoverVtErrorCode(err) != vtrpcpb.ErrorCode_CANCELLED {
errorCodeCheck = fmt.Errorf("Block command should only error due to canceled context: %v", err)
}
// Stream has finished.
break
}
if !firstLineReceived {
firstLineReceived = true
// The first log line will come from the "Block" command, so we are sure
// now that vtworker is actually executing it.
close(blockCommandStarted)
}
}
wg.Done()
}()
// Try to run a second, concurrent vtworker command.
// vtworker should send an error back that it's busy and we should retry later.
<-blockCommandStarted
gotErr := runVtworkerCommand(client, []string{"Ping", "Are you busy?"})
wantCode := vtrpcpb.ErrorCode_TRANSIENT_ERROR
if gotCode := vterrors.RecoverVtErrorCode(gotErr); gotCode != wantCode {
t.Fatalf("wrong error code for second cmd: got = %v, want = %v, err: %v", gotCode, wantCode, gotErr)
}
// Cancel running "Block" command.
if serverSideCancelation {
if err := runVtworkerCommand(client, []string{"Cancel"}); err != nil {
t.Fatal(err)
}
}
// Always cancel the context to not leak it (regardless of client or server
// side cancelation).
cancel()
wg.Wait()
if errorCodeCheck != nil {
t.Fatalf("Block command did not return the CANCELED error code: %v", errorCodeCheck)
}
// vtworker is now in a special state where the current job is already
// canceled but not reset yet. New commands are still failing with a
// retryable error.
gotErr2 := runVtworkerCommand(client, []string{"Ping", "canceled and still busy?"})
wantCode2 := vtrpcpb.ErrorCode_TRANSIENT_ERROR
if gotCode2 := vterrors.RecoverVtErrorCode(gotErr2); gotCode2 != wantCode2 {
t.Fatalf("wrong error code for second cmd before reset: got = %v, want = %v, err: %v", gotCode2, wantCode2, gotErr2)
}
// Reset vtworker for the next test function.
if err := resetVtworker(t, client); err != nil {
t.Fatal(err)
}
// Second vtworker command should succeed now after the first has finished.
if err := runVtworkerCommand(client, []string{"Ping", "You should not be busy anymore!"}); err != nil {
t.Fatalf("second cmd should not have failed: %v", err)
}
// Reset vtworker for the next test function.
if err := runVtworkerCommand(client, []string{"Reset"}); err != nil {
t.Fatal(err)
}
}