github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/cfg/dynamic_test.go (about)

     1  package cfg
     2  
     3  import (
     4  	"flag"
     5  	"io/ioutil"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/grafana/dskit/flagext"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func Test_DynamicUnmarshal(t *testing.T) {
    15  	defaultYamlConfig := `---
    16  server:
    17    port: 8080
    18  `
    19  
    20  	testContext := func(mockApplyDynamicConfig Source, config string, args []string) DynamicConfig {
    21  		data := NewDynamicConfig(mockApplyDynamicConfig)
    22  		fs := flag.NewFlagSet(t.Name(), flag.PanicOnError)
    23  
    24  		file, err := ioutil.TempFile("", "config.yaml")
    25  		require.NoError(t, err)
    26  		_, err = file.WriteString(config)
    27  		require.NoError(t, err)
    28  
    29  		configFileArgs := []string{"-config.file", file.Name()}
    30  		if args == nil {
    31  			args = configFileArgs
    32  		} else {
    33  			args = append(args, configFileArgs...)
    34  		}
    35  
    36  		err = DynamicUnmarshal(&data, args, fs)
    37  		require.NoError(t, err)
    38  		return data
    39  	}
    40  
    41  	t.Run("parses defaults", func(t *testing.T) {
    42  		data := testContext(nil, "", nil)
    43  
    44  		assert.Equal(t, 80, data.Server.Port)
    45  		assert.Equal(t, 60*time.Second, data.Server.Timeout)
    46  	})
    47  
    48  	t.Run("parses config from config.file", func(t *testing.T) {
    49  		data := testContext(nil, defaultYamlConfig, nil)
    50  		assert.Equal(t, 8080, data.Server.Port)
    51  	})
    52  
    53  	t.Run("calls ApplyDynamicConfig on provided DynamicCloneable", func(t *testing.T) {
    54  		applyDynamicConfigCalled := false
    55  		mockApplyDynamicConfig := func(dst Cloneable) error {
    56  			applyDynamicConfigCalled = true
    57  			return nil
    58  		}
    59  		data := testContext(mockApplyDynamicConfig, "", nil)
    60  		assert.NotNil(t, data)
    61  		assert.True(t, applyDynamicConfigCalled)
    62  	})
    63  
    64  	t.Run("makes config from file available to ApplyDynamicConfig", func(t *testing.T) {
    65  		var configFromFile *DynamicConfig
    66  		mockApplyDynamicConfig := func(dst Cloneable) error {
    67  			var ok bool
    68  			configFromFile, ok = dst.(*DynamicConfig)
    69  			require.True(t, ok)
    70  			return nil
    71  		}
    72  
    73  		data := testContext(mockApplyDynamicConfig, defaultYamlConfig, nil)
    74  		assert.NotNil(t, data)
    75  		assert.NotNil(t, configFromFile)
    76  		assert.Equal(t, 8080, configFromFile.Server.Port)
    77  	})
    78  
    79  	t.Run("config from file take precedence over config applied in ApplyDynamicConfig", func(t *testing.T) {
    80  		mockApplyDynamicConfig := func(dst Cloneable) error {
    81  			config, ok := dst.(*DynamicConfig)
    82  			require.True(t, ok)
    83  			config.Server.Port = 9090
    84  			return nil
    85  		}
    86  
    87  		data := testContext(mockApplyDynamicConfig, defaultYamlConfig, nil)
    88  		assert.Equal(t, 8080, data.Server.Port)
    89  	})
    90  
    91  	t.Run("config from command line takes precedence over config applied in ApplyDynamicConfig and in file", func(t *testing.T) {
    92  		mockApplyDynamicConfig := func(dst Cloneable) error {
    93  			config, ok := dst.(*DynamicConfig)
    94  			require.True(t, ok)
    95  			config.Server.Port = 9090
    96  			return nil
    97  		}
    98  
    99  		args := []string{
   100  			"-server.port", "7070",
   101  		}
   102  
   103  		data := testContext(mockApplyDynamicConfig, defaultYamlConfig, args)
   104  		assert.Equal(t, 7070, data.Server.Port)
   105  	})
   106  }
   107  
   108  type DynamicConfig struct {
   109  	Server             `yaml:"server"`
   110  	ConfigFile         string
   111  	applyDynamicConfig Source
   112  }
   113  
   114  func NewDynamicConfig(applyDynamicConfig Source) DynamicConfig {
   115  	if applyDynamicConfig == nil {
   116  		applyDynamicConfig = func(config Cloneable) error {
   117  			return nil
   118  		}
   119  	}
   120  	return DynamicConfig{
   121  		Server:             Server{},
   122  		ConfigFile:         "",
   123  		applyDynamicConfig: applyDynamicConfig,
   124  	}
   125  }
   126  
   127  func (d *DynamicConfig) Clone() flagext.Registerer {
   128  	return func(d DynamicConfig) *DynamicConfig {
   129  		return &d
   130  	}(*d)
   131  }
   132  
   133  func (d *DynamicConfig) ApplyDynamicConfig() Source {
   134  	return d.applyDynamicConfig
   135  }
   136  
   137  func (d *DynamicConfig) RegisterFlags(fs *flag.FlagSet) {
   138  	fs.IntVar(&d.Server.Port, "server.port", 80, "")
   139  	fs.DurationVar(&d.Server.Timeout, "server.timeout", 60*time.Second, "")
   140  	fs.StringVar(&d.ConfigFile, "config.file", "", "yaml file to load")
   141  }