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  }