github.com/matrixorigin/matrixone@v1.2.0/pkg/util/cn_address_func_test.go (about) 1 // Copyright 2022 Matrix Origin 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 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package util 16 17 import ( 18 "context" 19 "sync" 20 "testing" 21 22 "github.com/matrixorigin/matrixone/pkg/common/runtime" 23 "github.com/matrixorigin/matrixone/pkg/logutil" 24 log "github.com/matrixorigin/matrixone/pkg/pb/logservice" 25 "github.com/matrixorigin/matrixone/pkg/pb/metadata" 26 "github.com/stretchr/testify/assert" 27 ) 28 29 type testHAKeeperClient struct { 30 sync.RWMutex 31 value log.ClusterDetails 32 } 33 34 func (c *testHAKeeperClient) addCN(state metadata.WorkState, serviceIDs ...string) { 35 c.Lock() 36 defer c.Unlock() 37 for _, id := range serviceIDs { 38 c.value.CNStores = append(c.value.CNStores, log.CNStore{ 39 UUID: id, 40 SQLAddress: id, 41 WorkState: state, 42 }) 43 } 44 } 45 46 func (c *testHAKeeperClient) GetClusterDetails(ctx context.Context) (log.ClusterDetails, error) { 47 c.Lock() 48 defer c.Unlock() 49 return c.value, nil 50 } 51 52 func TestAddressFunc(t *testing.T) { 53 runtime.SetupProcessLevelRuntime(runtime.NewRuntime(metadata.ServiceType_CN, "test", logutil.GetGlobalLogger())) 54 t.Run("no client", func(t *testing.T) { 55 getClient := func() HAKeeperClient { 56 return nil 57 } 58 fn := AddressFunc(getClient) 59 ctx := context.Background() 60 cn, err := fn(ctx, true) 61 assert.Error(t, err) 62 assert.Equal(t, "", cn) 63 }) 64 65 t.Run("no cn", func(t *testing.T) { 66 getClient := func() HAKeeperClient { 67 return &testHAKeeperClient{} 68 } 69 fn := AddressFunc(getClient) 70 ctx := context.Background() 71 cn, err := fn(ctx, true) 72 assert.Error(t, err) 73 assert.Equal(t, "", cn) 74 }) 75 76 t.Run("one available cn", func(t *testing.T) { 77 client := &testHAKeeperClient{} 78 client.addCN(metadata.WorkState_Working, "cn1") 79 client.addCN(metadata.WorkState_Draining, "cn2") 80 getClient := func() HAKeeperClient { 81 return client 82 } 83 fn := AddressFunc(getClient) 84 ctx := context.Background() 85 cn, err := fn(ctx, false) // set false to return the last one. 86 assert.NoError(t, err) 87 assert.Equal(t, "cn1", cn) 88 }) 89 90 t.Run("multi cn", func(t *testing.T) { 91 client := &testHAKeeperClient{} 92 client.addCN(metadata.WorkState_Working, "cn1") 93 client.addCN(metadata.WorkState_Working, "cn2") 94 client.addCN(metadata.WorkState_Working, "cn3") 95 client.addCN(metadata.WorkState_Working, "cn4") 96 getClient := func() HAKeeperClient { 97 return client 98 } 99 fn := AddressFunc(getClient) 100 ctx := context.Background() 101 _, err := fn(ctx, true) 102 assert.NoError(t, err) 103 }) 104 }