github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/jobmaster/dm/bootstrap/upgrade_test.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package bootstrap
    15  
    16  import (
    17  	"context"
    18  	"sync"
    19  	"testing"
    20  
    21  	"github.com/coreos/go-semver/semver"
    22  	"github.com/pingcap/log"
    23  	"github.com/pingcap/tiflow/pkg/errors"
    24  	"github.com/stretchr/testify/mock"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  func TestUpgrade(t *testing.T) {
    29  	var (
    30  		dummyUpgrader = NewDummyUpgrader()
    31  		fromVer       = semver.New("6.1.0")
    32  		ver           = "6.2.0"
    33  	)
    34  
    35  	dummyUpgrader.On("UpgradeFuncs").Return([]UpgradeFunc{}).Once()
    36  	require.NoError(t, dummyUpgrader.Upgrade(context.Background(), *fromVer))
    37  
    38  	dummyUpgrader.On("UpgradeFuncs").Return([]UpgradeFunc{
    39  		{
    40  			Version: *semver.New("6.1.0"),
    41  			Upgrade: func(ctx context.Context) error {
    42  				ver = "6.1.0"
    43  				return nil
    44  			},
    45  		},
    46  		{
    47  			Version: *semver.New("6.3.0"),
    48  			Upgrade: func(ctx context.Context) error {
    49  				ver = "6.3.0"
    50  				return nil
    51  			},
    52  		},
    53  		{
    54  			Version: *semver.New("6.2.1"),
    55  			Upgrade: func(ctx context.Context) error {
    56  				ver = "6.2.1"
    57  				return nil
    58  			},
    59  		},
    60  	}).Once()
    61  	require.NoError(t, dummyUpgrader.Upgrade(context.Background(), *fromVer))
    62  	require.Equal(t, "6.3.0", ver)
    63  
    64  	dummyUpgrader.On("UpgradeFuncs").Return([]UpgradeFunc{
    65  		{
    66  			Version: *semver.New("6.3.0"),
    67  			Upgrade: func(ctx context.Context) error {
    68  				ver = "6.3.0"
    69  				return errors.New("failed to upgrade to v6.3.0")
    70  			},
    71  		},
    72  		{
    73  			Version: *semver.New("6.2.1"),
    74  			Upgrade: func(ctx context.Context) error {
    75  				ver = "6.2.1"
    76  				return nil
    77  			},
    78  		},
    79  	}).Once()
    80  	require.EqualError(t, dummyUpgrader.Upgrade(context.Background(), *fromVer), "failed to upgrade to v6.3.0")
    81  	require.Equal(t, "6.3.0", ver)
    82  
    83  	dummyUpgrader.On("UpgradeFuncs").Return([]UpgradeFunc{
    84  		{
    85  			Version: *semver.New("6.3.0"),
    86  			Upgrade: func(ctx context.Context) error {
    87  				ver = "6.3.0"
    88  				return errors.New("failed to upgrade to v6.3.0")
    89  			},
    90  			Rollback: func(ctx context.Context) error {
    91  				ver = "6.2.1"
    92  				return errors.New("failed to rollback to v6.2.1")
    93  			},
    94  		},
    95  		{
    96  			Version: *semver.New("6.2.1"),
    97  			Upgrade: func(ctx context.Context) error {
    98  				ver = "6.2.1"
    99  				return nil
   100  			},
   101  			Rollback: func(ctx context.Context) error {
   102  				ver = "6.2.0"
   103  				return errors.New("failed to rollback to v6.2.0")
   104  			},
   105  		},
   106  	}).Once()
   107  	require.EqualError(t, dummyUpgrader.Upgrade(context.Background(), *fromVer), "failed to upgrade to v6.3.0")
   108  	require.Equal(t, "6.2.0", ver)
   109  }
   110  
   111  type DummyUpgrader struct {
   112  	*DefaultUpgrader
   113  
   114  	mock.Mock
   115  	mu sync.Mutex
   116  }
   117  
   118  func NewDummyUpgrader() *DummyUpgrader {
   119  	u := &DummyUpgrader{
   120  		DefaultUpgrader: NewDefaultUpgrader(log.L()),
   121  	}
   122  	u.DefaultUpgrader.Upgrader = u
   123  	return u
   124  }
   125  
   126  func (upgrader *DummyUpgrader) UpgradeFuncs() []UpgradeFunc {
   127  	upgrader.mu.Lock()
   128  	defer upgrader.mu.Unlock()
   129  	args := upgrader.Called()
   130  	return args.Get(0).([]UpgradeFunc)
   131  }