From 981de21183ad935113f2abc6d610cc948186f6b8 Mon Sep 17 00:00:00 2001 From: Daniel Einspanjer Date: Wed, 25 Jul 2018 09:23:24 -0400 Subject: [PATCH] Issue #539 - Add viper.CancelWatchConfig() --- viper.go | 18 +++++++-- viper_test.go | 103 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 3 deletions(-) diff --git a/viper.go b/viper.go index 907a102..d8c78e6 100644 --- a/viper.go +++ b/viper.go @@ -184,6 +184,7 @@ type Viper struct { properties *properties.Properties onConfigChange func(fsnotify.Event) + watchChannel chan bool } // New returns an initialized Viper instance. @@ -277,7 +278,8 @@ func (v *Viper) WatchConfig() { configFile := filepath.Clean(filename) configDir, _ := filepath.Split(configFile) - done := make(chan bool) + v.watchChannel = make(chan bool) + go func() { for { select { @@ -293,16 +295,26 @@ func (v *Viper) WatchConfig() { } } case err := <-watcher.Errors: - log.Println("error:", err) + if err != nil { + log.Println("error:", err) + } } } }() watcher.Add(configDir) - <-done + <-v.watchChannel + watcher.Close() }() } +func CancelWatchConfig() { + v.CancelWatchConfig() +} +func (v *Viper) CancelWatchConfig() { + v.watchChannel <- true +} + // SetConfigFile explicitly defines the path, name and extension of the config file. // Viper will use this and not check any of the config paths. func SetConfigFile(in string) { v.SetConfigFile(in) } diff --git a/viper_test.go b/viper_test.go index 60543f5..06ba577 100644 --- a/viper_test.go +++ b/viper_test.go @@ -7,6 +7,7 @@ package viper import ( "bytes" + "context" "fmt" "io" "io/ioutil" @@ -18,6 +19,7 @@ import ( "testing" "time" + "github.com/fsnotify/fsnotify" "github.com/spf13/afero" "github.com/spf13/cast" @@ -768,6 +770,107 @@ func TestIsSet(t *testing.T) { assert.True(t, v.IsSet("helloworld")) } +func TestWatchConfig(t *testing.T) { + before := []byte("key = \"value is before\"\n") + changed := []byte("key = \"value is changed\"\n") + after := []byte("key = \"value is after\"\n") + + // This message is used for true asserts just making sure the file changes are happening as expected. + errMsg := "Test threw an unexpected error (test error)" + + // Context is used for timeout handling within test. + var ( + ctx context.Context + cancel context.CancelFunc + ) + + // Get a reference to a temporary file that we'll use for the test. Note we can't set a file extension using this API. + tmpfile, err := ioutil.TempFile("", "watch-test") + assert.Nil(t, err, errMsg) + defer os.Remove(tmpfile.Name()) // clean up + + t.Log("Writing initial value to test config file: ", tmpfile.Name()) + _, err = tmpfile.Write(before) + assert.Nil(t, err, errMsg) + + err = tmpfile.Close() + assert.Nil(t, err, errMsg) + + // This block just confirms that we properly wrote the desired contents to the file. + data, err := ioutil.ReadFile(tmpfile.Name()) + assert.Nil(t, err, errMsg) + assert.Equal(t, data, before, "Contents of test file are unexpected (test error)") + + // Set up a viper to test backing it with the tmp config file + v := New() + v.SetDefault(`key`, `value is default`) + v.SetConfigFile(tmpfile.Name()) + v.SetConfigType("toml") + + t.Log("Initial read of config") + err = v.ReadInConfig() + assert.Nil(t, err, errMsg) + + assert.Equal(t, `value is before`, v.GetString(`key`), "viper did not see the correct initial value for the config file.") + + // Set up a context with deadline so we won't wait forever if we don't see a change. + ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + // Turn on watch config. + v.WatchConfig() + v.OnConfigChange(func(e fsnotify.Event) { + t.Log("OnConfigChange executed.") + // If we get the signal of a change, we can immediately cancel the context. + cancel() + }) + + // XXX: I'm not sure why, but if we don't sleep, the watcher won't see a change that happens immediately after WatchConfig() is called. + time.Sleep(50 * time.Millisecond) + t.Log("Writing changed value to ", tmpfile.Name()) + + // Not sure if the TRUNC is necessary, but basically re-open the file so we can change it. + tmpfile, err = os.OpenFile(tmpfile.Name(), os.O_TRUNC|os.O_RDWR, 0644) + assert.Nil(t, err, errMsg) + + _, err = tmpfile.Write(changed) + assert.Nil(t, err, errMsg) + + err = tmpfile.Close() + assert.Nil(t, err, errMsg) + + data, err = ioutil.ReadFile(tmpfile.Name()) + assert.Nil(t, err, errMsg) + assert.Equal(t, data, changed, "Contents of test file are unexpected (test error)") + + // Wait here for either a timeout or signal that we saw a change. + <-ctx.Done() + assert.NotEqual(t, context.DeadlineExceeded, ctx.Err(), "Timed out waiting for change notification.") + + assert.Equal(t, `value is changed`, v.GetString(`key`), "viper did not see the correct changed value for the config file after WatchConfig().") + + // Canceling turns off the fsevent Watcher so even if we end up starting a new viper instance (or calling viper.Reset()) we won't pick up spurious change events. + v.CancelWatchConfig() + // XXX: I'm not sure why, but if we don't sleep, the watcher might not fully close down before the next event happens. Doesn't affect this test, but can cause an error on the next one. + time.Sleep(50 * time.Millisecond) + + // Now we make one more change to the file to verify the viper config doesn't pick up the change. + tmpfile, err = os.OpenFile(tmpfile.Name(), os.O_TRUNC|os.O_RDWR, 0644) + assert.Nil(t, err, errMsg) + + _, err = tmpfile.Write(after) + assert.Nil(t, err, errMsg) + + err = tmpfile.Close() + assert.Nil(t, err, errMsg) + + data, err = ioutil.ReadFile(tmpfile.Name()) + assert.Nil(t, err, errMsg) + assert.Equal(t, data, after, "Contents of test file are unexpected (test error)") + + assert.NotEqual(t, `value is after`, v.GetString(`key`), "viper saw the after value in the file even after StopWatchConfig().") +} + func TestDirsSearch(t *testing.T) { root, config, cleanup := initDirs(t)