github.com/pingcap/ticdc@v0.0.0-20220526033649-485a10ef2652/pkg/orchestrator/reactor_state_tester.go (about)

     1  // Copyright 2021 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 orchestrator
    15  
    16  import (
    17  	"github.com/pingcap/check"
    18  	"github.com/pingcap/errors"
    19  	cerrors "github.com/pingcap/ticdc/pkg/errors"
    20  	"github.com/pingcap/ticdc/pkg/orchestrator/util"
    21  )
    22  
    23  // ReactorStateTester is a helper struct for unit-testing an implementer of ReactorState
    24  type ReactorStateTester struct {
    25  	c         *check.C
    26  	state     ReactorState
    27  	kvEntries map[string]string
    28  }
    29  
    30  // NewReactorStateTester creates a new ReactorStateTester
    31  func NewReactorStateTester(c *check.C, state ReactorState, initKVEntries map[string]string) *ReactorStateTester {
    32  	if initKVEntries == nil {
    33  		initKVEntries = make(map[string]string)
    34  	}
    35  	for k, v := range initKVEntries {
    36  		err := state.Update(util.NewEtcdKey(k), []byte(v), true)
    37  		c.Assert(err, check.IsNil)
    38  	}
    39  	return &ReactorStateTester{
    40  		c:         c,
    41  		state:     state,
    42  		kvEntries: initKVEntries,
    43  	}
    44  }
    45  
    46  // Update is used to update keys in the mocked kv-store.
    47  func (t *ReactorStateTester) Update(key string, value []byte) error {
    48  	k := util.NewEtcdKey(key)
    49  	err := t.state.Update(k, value, false)
    50  	if err != nil {
    51  		return errors.Trace(err)
    52  	}
    53  	if value != nil {
    54  		t.kvEntries[key] = string(value)
    55  	} else {
    56  		delete(t.kvEntries, key)
    57  	}
    58  	return nil
    59  }
    60  
    61  // ApplyPatches calls the GetPatches method on the ReactorState and apply the changes to the mocked kv-store.
    62  func (t *ReactorStateTester) ApplyPatches() error {
    63  	patchGroups := t.state.GetPatches()
    64  	for _, patches := range patchGroups {
    65  		err := t.applyPatches(patches)
    66  		if err != nil {
    67  			return err
    68  		}
    69  	}
    70  	return nil
    71  }
    72  
    73  func (t *ReactorStateTester) applyPatches(patches []DataPatch) error {
    74  RetryLoop:
    75  	for {
    76  		tmpKVEntries := make(map[util.EtcdKey][]byte)
    77  		for k, v := range t.kvEntries {
    78  			tmpKVEntries[util.NewEtcdKey(k)] = []byte(v)
    79  		}
    80  		changedSet := make(map[util.EtcdKey]struct{})
    81  		for _, patch := range patches {
    82  			err := patch.Patch(tmpKVEntries, changedSet)
    83  			if cerrors.ErrEtcdIgnore.Equal(errors.Cause(err)) {
    84  				continue
    85  			} else if cerrors.ErrEtcdTryAgain.Equal(errors.Cause(err)) {
    86  				continue RetryLoop
    87  			} else if err != nil {
    88  				return errors.Trace(err)
    89  			}
    90  		}
    91  		for k := range changedSet {
    92  			err := t.state.Update(k, tmpKVEntries[k], false)
    93  			if err != nil {
    94  				return err
    95  			}
    96  			if value := tmpKVEntries[k]; value != nil {
    97  				t.kvEntries[k.String()] = string(value)
    98  			} else {
    99  				delete(t.kvEntries, k.String())
   100  			}
   101  		}
   102  		return nil
   103  	}
   104  }
   105  
   106  // MustApplyPatches calls ApplyPatches and must successfully
   107  func (t *ReactorStateTester) MustApplyPatches() {
   108  	t.c.Assert(t.ApplyPatches(), check.IsNil)
   109  }
   110  
   111  // MustUpdate calls Update and must successfully
   112  func (t *ReactorStateTester) MustUpdate(key string, value []byte) {
   113  	t.c.Assert(t.Update(key, value), check.IsNil)
   114  }
   115  
   116  // KVEntries returns the contents of the mocked KV store.
   117  func (t *ReactorStateTester) KVEntries() map[string]string {
   118  	return t.kvEntries
   119  }