github.com/go-spring/spring-base@v1.1.3/run/run.go (about)

     1  /*
     2   * Copyright 2012-2019 the original author or authors.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *      https://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  // Package run provides methods to query the running mode of program.
    18  package run
    19  
    20  import (
    21  	"errors"
    22  	"os"
    23  	"strings"
    24  )
    25  
    26  var mode int
    27  
    28  const (
    29  	NormalModeFlag = 0x0000
    30  	RecordModeFlag = 0x0001
    31  	ReplayModeFlag = 0x0002
    32  	TestModeFlag   = 0x0004
    33  )
    34  
    35  func init() {
    36  	initMode()
    37  }
    38  
    39  // initMode for unit test.
    40  func initMode() {
    41  	mode = NormalModeFlag
    42  	if os.Getenv("GS_RECORD_MODE") != "" {
    43  		mode |= RecordModeFlag
    44  	}
    45  	if os.Getenv("GS_REPLAY_MODE") != "" {
    46  		mode |= ReplayModeFlag
    47  	}
    48  	for _, arg := range os.Args {
    49  		if strings.HasPrefix(arg, "-test.") {
    50  			mode |= TestModeFlag
    51  			break
    52  		}
    53  	}
    54  }
    55  
    56  // SetMode sets the running mode, only in unit test mode.
    57  func SetMode(flag int) (reset func()) {
    58  	MustTestMode()
    59  	old := mode
    60  	reset = func() {
    61  		mode = old
    62  	}
    63  	mode = flag
    64  	return
    65  }
    66  
    67  // NormalMode returns whether it is running in normal mode.
    68  func NormalMode() bool {
    69  	return mode == 0
    70  }
    71  
    72  // RecordMode returns whether it is running in record mode.
    73  func RecordMode() bool {
    74  	return mode&RecordModeFlag == RecordModeFlag
    75  }
    76  
    77  // ReplayMode returns whether it is running in replay mode.
    78  func ReplayMode() bool {
    79  	return mode&ReplayModeFlag == ReplayModeFlag
    80  }
    81  
    82  // TestMode returns whether it is running in test mode.
    83  func TestMode() bool {
    84  	return mode&TestModeFlag == TestModeFlag
    85  }
    86  
    87  // MustTestMode panic occurs when calling not in test mode.
    88  func MustTestMode() {
    89  	if !TestMode() {
    90  		panic(errors.New("should be called in test mode"))
    91  	}
    92  }