vitess.io/vitess@v0.16.2/go/vt/vtctl/workflow/server_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 workflow
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"testing"
    23  
    24  	"github.com/stretchr/testify/assert"
    25  	"github.com/stretchr/testify/require"
    26  	"google.golang.org/protobuf/encoding/prototext"
    27  
    28  	"vitess.io/vitess/go/sqltypes"
    29  	"vitess.io/vitess/go/test/utils"
    30  	"vitess.io/vitess/go/vt/topo/topoproto"
    31  	"vitess.io/vitess/go/vt/vttablet/tmclient"
    32  
    33  	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
    34  	querypb "vitess.io/vitess/go/vt/proto/query"
    35  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    36  )
    37  
    38  type fakeTMC struct {
    39  	tmclient.TabletManagerClient
    40  	vrepQueriesByTablet map[string]map[string]*querypb.QueryResult
    41  }
    42  
    43  func (fake *fakeTMC) VReplicationExec(ctx context.Context, tablet *topodatapb.Tablet, query string) (*querypb.QueryResult, error) {
    44  	alias := topoproto.TabletAliasString(tablet.Alias)
    45  	tabletQueries, ok := fake.vrepQueriesByTablet[alias]
    46  	if !ok {
    47  		return nil, fmt.Errorf("no query map registered on fake for %s", alias)
    48  	}
    49  
    50  	p3qr, ok := tabletQueries[query]
    51  	if !ok {
    52  		return nil, fmt.Errorf("no result on fake for query %q on tablet %s", query, alias)
    53  	}
    54  
    55  	return p3qr, nil
    56  }
    57  
    58  func TestCheckReshardingJournalExistsOnTablet(t *testing.T) {
    59  	t.Parallel()
    60  
    61  	ctx := context.Background()
    62  	tablet := &topodatapb.Tablet{
    63  		Alias: &topodatapb.TabletAlias{
    64  			Cell: "zone1",
    65  			Uid:  100,
    66  		},
    67  	}
    68  	journal := &binlogdatapb.Journal{
    69  		Id:            1,
    70  		MigrationType: binlogdatapb.MigrationType_SHARDS,
    71  		Tables:        []string{"t1", "t2"},
    72  	}
    73  	journalBytes, err := prototext.Marshal(journal)
    74  	require.NoError(t, err, "could not marshal journal %+v into bytes", journal)
    75  
    76  	// get some bytes that will fail to unmarshal into a binlogdatapb.Journal
    77  	tabletBytes, err := prototext.Marshal(tablet)
    78  	require.NoError(t, err, "could not marshal tablet %+v into bytes", tablet)
    79  
    80  	p3qr := sqltypes.ResultToProto3(sqltypes.MakeTestResult([]*querypb.Field{
    81  		{
    82  			Name: "val",
    83  			Type: querypb.Type_BLOB,
    84  		},
    85  	}, string(journalBytes)))
    86  
    87  	tests := []struct {
    88  		name        string
    89  		tablet      *topodatapb.Tablet
    90  		result      *querypb.QueryResult
    91  		journal     *binlogdatapb.Journal
    92  		shouldExist bool
    93  		shouldErr   bool
    94  	}{
    95  		{
    96  			name:        "journal exists",
    97  			tablet:      tablet,
    98  			result:      p3qr,
    99  			shouldExist: true,
   100  			journal:     journal,
   101  		},
   102  		{
   103  			name:        "journal does not exist",
   104  			tablet:      tablet,
   105  			result:      sqltypes.ResultToProto3(sqltypes.MakeTestResult(nil)),
   106  			journal:     &binlogdatapb.Journal{},
   107  			shouldExist: false,
   108  		},
   109  		{
   110  			name:   "cannot unmarshal into journal",
   111  			tablet: tablet,
   112  			result: sqltypes.ResultToProto3(sqltypes.MakeTestResult([]*querypb.Field{
   113  				{
   114  					Name: "val",
   115  					Type: querypb.Type_BLOB,
   116  				},
   117  			}, string(tabletBytes))),
   118  			shouldErr: true,
   119  		},
   120  		{
   121  			name: "VReplicationExec fails on tablet",
   122  			tablet: &topodatapb.Tablet{ // Here we use a different tablet to force the fake to return an error
   123  				Alias: &topodatapb.TabletAlias{
   124  					Cell: "zone2",
   125  					Uid:  200,
   126  				},
   127  			},
   128  			shouldErr: true,
   129  		},
   130  	}
   131  
   132  	for _, tt := range tests {
   133  		tt := tt
   134  		t.Run(tt.name, func(t *testing.T) {
   135  			t.Parallel()
   136  
   137  			tmc := &fakeTMC{
   138  				vrepQueriesByTablet: map[string]map[string]*querypb.QueryResult{
   139  					topoproto.TabletAliasString(tablet.Alias): { // always use the tablet shared by these tests cases
   140  						"select val from _vt.resharding_journal where id=1": tt.result,
   141  					},
   142  				},
   143  			}
   144  
   145  			ws := NewServer(nil, tmc)
   146  			journal, exists, err := ws.CheckReshardingJournalExistsOnTablet(ctx, tt.tablet, 1)
   147  			if tt.shouldErr {
   148  				assert.Error(t, err)
   149  				return
   150  			}
   151  
   152  			require.NoError(t, err)
   153  
   154  			existAssertionMsg := "expected journal to "
   155  			if tt.shouldExist {
   156  				existAssertionMsg += "already exist on tablet"
   157  			} else {
   158  				existAssertionMsg += "not exist"
   159  			}
   160  
   161  			assert.Equal(t, tt.shouldExist, exists, existAssertionMsg)
   162  			utils.MustMatch(t, tt.journal, journal, "journal in resharding_journal did not match")
   163  		})
   164  	}
   165  }