From a7084797945bd6ddcca91e4981b82c7393e6dca1 Mon Sep 17 00:00:00 2001
From: javaducky <javaducky@gmail.com>
Date: Fri, 13 Sep 2019 09:29:19 -0500
Subject: [PATCH] Updating SafeWriteConfig and SafeWriteConfigAs to match
 documented behavior.

 Methods should throw an error if the config file already exists or if no configpath is configured when not explicitly requesting a write path.
---
 viper.go      | 27 ++++++++++++++---
 viper_test.go | 80 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/viper.go b/viper.go
index 4c98c50..6167aff 100644
--- a/viper.go
+++ b/viper.go
@@ -115,6 +115,22 @@ func (fnfe ConfigFileNotFoundError) Error() string {
 	return fmt.Sprintf("Config File %q Not Found in %q", fnfe.name, fnfe.locations)
 }
 
+// ConfigFileAlreadyExistsError denotes failure to write new configuration file.
+type ConfigFileAlreadyExistsError string
+
+// Error returns the formatted error when configuration already exists.
+func (faee ConfigFileAlreadyExistsError) Error() string {
+	return fmt.Sprintf("Config File %q Already Exists", string(faee))
+}
+
+// MissingConfigurationError denotes a required configuration setting has not been provided.
+type MissingConfigurationError string
+
+// Error returns the formatted error when a required configuration element has not been provided.
+func (mce MissingConfigurationError) Error() string {
+	return fmt.Sprintf("Missing Configuration for %q", string(mce))
+}
+
 // A DecoderConfigOption can be passed to viper.Unmarshal to configure
 // mapstructure.DecoderConfig options
 type DecoderConfigOption func(*mapstructure.DecoderConfig)
@@ -1337,11 +1353,10 @@ func (v *Viper) WriteConfig() error {
 // SafeWriteConfig writes current configuration to file only if the file does not exist.
 func SafeWriteConfig() error { return v.SafeWriteConfig() }
 func (v *Viper) SafeWriteConfig() error {
-	filename, err := v.getConfigFile()
-	if err != nil {
-		return err
+	if len(v.configPaths) < 1 {
+		return MissingConfigurationError("configPath")
 	}
-	return v.writeConfig(filename, false)
+	return v.SafeWriteConfigAs(filepath.Join(v.configPaths[0], v.configName+"."+v.configType))
 }
 
 // WriteConfigAs writes current configuration to a given filename.
@@ -1353,6 +1368,10 @@ func (v *Viper) WriteConfigAs(filename string) error {
 // SafeWriteConfigAs writes current configuration to a given filename if it does not exist.
 func SafeWriteConfigAs(filename string) error { return v.SafeWriteConfigAs(filename) }
 func (v *Viper) SafeWriteConfigAs(filename string) error {
+	handle, err := v.fs.Stat(filename)
+	if handle != nil && err == nil {
+		return ConfigFileAlreadyExistsError(filename)
+	}
 	return v.writeConfig(filename, false)
 }
 
diff --git a/viper_test.go b/viper_test.go
index 677a14b..7651787 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -1384,6 +1384,86 @@ func TestWriteConfigYAML(t *testing.T) {
 	assert.Equal(t, yamlWriteExpected, read)
 }
 
+func TestSafeWriteConfig(t *testing.T) {
+	v := New()
+	fs := afero.NewMemMapFs()
+	v.SetFs(fs)
+	v.AddConfigPath("/test")
+	v.SetConfigName("c")
+	v.SetConfigType("yaml")
+	err := v.ReadConfig(bytes.NewBuffer(yamlExample))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = v.SafeWriteConfig(); err != nil {
+		t.Fatal(err)
+	}
+	read, err := afero.ReadFile(fs, "/test/c.yaml")
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, yamlWriteExpected, read)
+}
+
+func TestSafeWriteConfigWithMissingConfigPath(t *testing.T) {
+	v := New()
+	fs := afero.NewMemMapFs()
+	v.SetFs(fs)
+	v.SetConfigName("c")
+	v.SetConfigType("yaml")
+	err := v.SafeWriteConfig()
+	if err == nil {
+		t.Fatal("Expected exception")
+	}
+	_, ok := err.(MissingConfigurationError)
+	assert.True(t, ok, "Expected MissingConfigurationError")
+}
+
+func TestSafeWriteConfigWithExistingFile(t *testing.T) {
+	v := New()
+	fs := afero.NewMemMapFs()
+	fs.Create("/test/c.yaml")
+	v.SetFs(fs)
+	v.AddConfigPath("/test")
+	v.SetConfigName("c")
+	v.SetConfigType("yaml")
+	err := v.SafeWriteConfig()
+	if err == nil {
+		t.Fatal("Expected exception")
+	}
+	_, ok := err.(ConfigFileAlreadyExistsError)
+	assert.True(t, ok, "Expected ConfigFileAlreadyExistsError")
+}
+
+func TestSafeWriteAsConfig(t *testing.T) {
+	v := New()
+	fs := afero.NewMemMapFs()
+	v.SetFs(fs)
+	err := v.ReadConfig(bytes.NewBuffer(yamlExample))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = v.SafeWriteConfigAs("/test/c.yaml"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err = afero.ReadFile(fs, "/test/c.yaml"); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestSafeWriteConfigAsWithExistingFile(t *testing.T) {
+	v := New()
+	fs := afero.NewMemMapFs()
+	fs.Create("/test/c.yaml")
+	v.SetFs(fs)
+	err := v.SafeWriteConfigAs("/test/c.yaml")
+	if err == nil {
+		t.Fatal("Expected exception")
+	}
+	_, ok := err.(ConfigFileAlreadyExistsError)
+	assert.True(t, ok, "Expected ConfigFileAlreadyExistsError")
+}
+
 var yamlMergeExampleTgt = []byte(`
 hello:
     pop: 37890