github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/config/testutil.go (about) 1 // Copyright 2015 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package config 12 13 import ( 14 "github.com/cockroachdb/cockroach/pkg/config/zonepb" 15 "github.com/cockroachdb/cockroach/pkg/util/stop" 16 "github.com/cockroachdb/cockroach/pkg/util/syncutil" 17 ) 18 19 type zoneConfigMap map[uint32]zonepb.ZoneConfig 20 21 var ( 22 testingZoneConfig zoneConfigMap 23 testingHasHook bool 24 testingPreviousHook zoneConfigHook 25 testingLock syncutil.Mutex 26 ) 27 28 // TestingSetupZoneConfigHook initializes the zone config hook 29 // to 'testingZoneConfigHook' which uses 'testingZoneConfig'. 30 // Settings go back to their previous values when the stopper runs our closer. 31 func TestingSetupZoneConfigHook(stopper *stop.Stopper) { 32 stopper.AddCloser(stop.CloserFn(testingResetZoneConfigHook)) 33 34 testingLock.Lock() 35 defer testingLock.Unlock() 36 if testingHasHook { 37 panic("TestingSetupZoneConfigHook called without restoring state") 38 } 39 testingHasHook = true 40 testingZoneConfig = make(zoneConfigMap) 41 testingPreviousHook = ZoneConfigHook 42 ZoneConfigHook = testingZoneConfigHook 43 testingLargestIDHook = func(maxID uint32) (max uint32) { 44 testingLock.Lock() 45 defer testingLock.Unlock() 46 for id := range testingZoneConfig { 47 if maxID > 0 && id > maxID { 48 continue 49 } 50 if id > max { 51 max = id 52 } 53 } 54 return 55 } 56 } 57 58 // testingResetZoneConfigHook resets the zone config hook back to what it was 59 // before TestingSetupZoneConfigHook was called. 60 func testingResetZoneConfigHook() { 61 testingLock.Lock() 62 defer testingLock.Unlock() 63 if !testingHasHook { 64 panic("TestingResetZoneConfigHook called on uninitialized testing hook") 65 } 66 testingHasHook = false 67 ZoneConfigHook = testingPreviousHook 68 testingLargestIDHook = nil 69 } 70 71 // TestingSetZoneConfig sets the zone config entry for object 'id' 72 // in the testing map. 73 func TestingSetZoneConfig(id uint32, zone zonepb.ZoneConfig) { 74 testingLock.Lock() 75 defer testingLock.Unlock() 76 testingZoneConfig[id] = zone 77 } 78 79 func testingZoneConfigHook( 80 _ *SystemConfig, id uint32, 81 ) (*zonepb.ZoneConfig, *zonepb.ZoneConfig, bool, error) { 82 testingLock.Lock() 83 defer testingLock.Unlock() 84 if zone, ok := testingZoneConfig[id]; ok { 85 return &zone, nil, false, nil 86 } 87 return nil, nil, false, nil 88 }