Add yamledit package and preserve wordpress_sites.yml formatting

cmd/new.go

@@ -159,13 +159,13 @@ func (c *NewCommand) Run(args []string) int {
 	}
 
 	// Update default configs
-	for env, config := range c.trellis.Environments {
-		c.trellis.UpdateDefaultConfig(config, c.name, c.host, env)
-		_ = c.trellis.WriteYamlFile(
-			config,
-			filepath.Join("group_vars", env, "wordpress_sites.yml"),
-			c.YamlHeader("https://roots.io/trellis/docs/wordpress-sites/"),
-		)
+	for env := range c.trellis.Environments {
+		sitesPath := filepath.Join("group_vars", env, "wordpress_sites.yml")
+		if err := c.trellis.UpdateDefaultSiteFile(sitesPath, c.name, c.host, env, c.createdByComment()); err != nil {
+			c.UI.Error(fmt.Sprintf("Error updating %s: %s", sitesPath, err))
+			c.UI.Error("Your project may not work correctly.")
+			c.UI.Error("This is probably a trellis-cli bug. Please open an issue at: https://github.com/roots/trellis-cli")
+		}
 
 		stringGenerator := trellis.RandomStringGenerator{Length: 64}
 		vault := c.trellis.GenerateVaultConfig(c.name, env, &stringGenerator)
@@ -261,9 +261,11 @@ Options:
 }
 
 func (c *NewCommand) YamlHeader(doc string) string {
-	const header = "# Created by trellis-cli v%s\n# Documentation: %s\n\n"
+	return fmt.Sprintf("%s\n# Documentation: %s\n\n", c.createdByComment(), doc)
+}
 
-	return fmt.Sprintf(header, c.CliVersion, doc)
+func (c *NewCommand) createdByComment() string {
+	return fmt.Sprintf("# Created by trellis-cli v%s", c.CliVersion)
 }
 
 func askDomain(ui cli.Ui, path string) (host string, err error) {

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/posener/complete v1.2.3
 	github.com/theckman/yacspin v0.13.12
 	github.com/weppos/publicsuffix-go v0.50.3
+	go.yaml.in/yaml/v3 v3.0.5
 	golang.org/x/crypto v0.58.0
 	golang.org/x/oauth2 v0.37.0
 	gopkg.in/alessio/shellescape.v1 v1.0.0-20170105083845-52074bc9df61

trellis/config.go

@@ -6,8 +6,10 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
+	"strings"
 
 	"github.com/roots/trellis-cli/dns"
+	"github.com/roots/trellis-cli/trellis/yamledit"
 	"github.com/weppos/publicsuffix-go/publicsuffix"
 	"gopkg.in/yaml.v2"
 )
@@ -107,14 +109,52 @@ func (t *Trellis) HostsFromDomain(domain string, env string) (canonical *publics
 	}
 }
 
-func (t *Trellis) UpdateDefaultConfig(config *Config, name string, host string, env string) {
-	config.WordPressSites[name] = config.WordPressSites[DefaultSiteName]
+// UpdateDefaultSiteFile renames the default site in a wordpress_sites.yml file
+// (relative to the project) and applies the generated hosts and environment
+// defaults. Comments and site options not modeled by Site are preserved.
+func (t *Trellis) UpdateDefaultSiteFile(path string, name string, host string, env string, header string) error {
+	path = filepath.Join(t.Path, path)
 
-	if name != DefaultSiteName {
-		delete(config.WordPressSites, DefaultSiteName)
+	doc, err := yamledit.Load(path)
+	if err != nil {
+		return err
+	}
+
+	if err := doc.RenameKey(name, "wordpress_sites", DefaultSiteName); err != nil {
+		return err
+	}
+
+	site := &Site{}
+	if err := doc.Decode(site, "wordpress_sites", name); err != nil {
+		return err
+	}
+
+	t.GenerateSite(site, host, env)
+
+	if err := doc.Set(site.SiteHosts, "wordpress_sites", name, "site_hosts"); err != nil {
+		return err
+	}
+
+	optional := []struct{ key, value string }{
+		{"admin_email", site.AdminEmail},
+		{"branch", site.Branch},
+		{"repo", site.Repo},
+		{"repo_subtree_path", site.RepoSubtreePath},
+	}
+
+	for _, option := range optional {
+		if option.value == "" {
+			doc.Delete("wordpress_sites", name, option.key)
+		} else if err := doc.Set(option.value, "wordpress_sites", name, option.key); err != nil {
+			return err
+		}
+	}
+
+	if header != "" {
+		doc.SetHeadComment(strings.TrimRight(header+"\n"+doc.HeadComment(), "\n"))
 	}
 
-	t.GenerateSite(config.WordPressSites[name], host, env)
+	return doc.Save(path)
 }
 
 func (s *Site) SslEnabled() bool {

trellis/config_test.go

@@ -2,6 +2,8 @@ package trellis
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"testing"
@@ -10,7 +12,7 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
-func TestUpdateDefaultConfig(t *testing.T) {
+func TestUpdateDefaultSiteFile(t *testing.T) {
 	trellis := &Trellis{}
 
 	cases := []struct {
@@ -216,9 +218,14 @@ wordpress_sites:
 	}
 
 	for _, tc := range cases {
-		config := trellis.ParseConfig(fmt.Sprintf("testdata/trellis/group_vars/%s/wordpress_sites.yml", tc.env))
-		trellis.UpdateDefaultConfig(config, tc.siteName, tc.host, tc.env)
+		trellis.Path = t.TempDir()
+		copySitesFixture(t, trellis.Path, tc.env)
 
+		if err := trellis.UpdateDefaultSiteFile("wordpress_sites.yml", tc.siteName, tc.host, tc.env, ""); err != nil {
+			t.Fatalf("%s => unexpected error: %v", tc.name, err)
+		}
+
+		config := trellis.ParseConfig(filepath.Join(trellis.Path, "wordpress_sites.yml"))
 		configYaml, _ := yaml.Marshal(config)
 		expected := strings.TrimSpace(tc.expectedYaml)
 		output := strings.TrimSpace(string(configYaml))
@@ -229,6 +236,70 @@ wordpress_sites:
 	}
 }
 
+func TestUpdateDefaultSiteFilePreservesFormatting(t *testing.T) {
+	trellis := &Trellis{Path: t.TempDir()}
+	copySitesFixture(t, trellis.Path, "production")
+
+	if err := trellis.UpdateDefaultSiteFile("wordpress_sites.yml", "mysite.com", "mysite.com", "production", "# Created by trellis-cli v1.0.0"); err != nil {
+		t.Fatal(err)
+	}
+
+	output, _ := os.ReadFile(filepath.Join(trellis.Path, "wordpress_sites.yml"))
+
+	expected := `# Created by trellis-cli v1.0.0
+# Documentation: https://roots.io/trellis/docs/remote-server-setup/
+# ` + "`wordpress_sites`" + ` options: https://roots.io/trellis/docs/wordpress-sites
+# Define accompanying passwords/secrets in group_vars/production/vault.yml
+
+wordpress_sites:
+  mysite.com:
+    site_hosts:
+      - canonical: mysite.com
+        redirects:
+          - www.mysite.com
+    local_path: ../site # path targeting local Bedrock site directory (relative to Ansible root)
+    repo: git@github.com:example/example.com.git # replace with your Git repo URL
+    repo_subtree_path: site # relative path to your Bedrock/WP directory in your repo
+    branch: master
+    multisite:
+      enabled: false
+    ssl:
+      enabled: false
+      provider: letsencrypt
+    cache:
+      enabled: false
+`
+
+	if string(output) != expected {
+		t.Errorf("expected:\n%s\ngot:\n%s", expected, output)
+	}
+}
+
+func TestUpdateDefaultSiteFileMissingDefaultSite(t *testing.T) {
+	trellis := &Trellis{Path: t.TempDir()}
+
+	if err := os.WriteFile(filepath.Join(trellis.Path, "wordpress_sites.yml"), []byte("wordpress_sites:\n  other.com: {}\n"), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := trellis.UpdateDefaultSiteFile("wordpress_sites.yml", "mysite.com", "mysite.com", "production", ""); err == nil {
+		t.Error("expected error when default site is missing")
+	}
+}
+
+func copySitesFixture(t *testing.T, dir string, env string) {
+	t.Helper()
+
+	data, err := os.ReadFile(fmt.Sprintf("testdata/trellis/group_vars/%s/wordpress_sites.yml", env))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := os.WriteFile(filepath.Join(dir, "wordpress_sites.yml"), data, 0644); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestSslProvider(t *testing.T) {
 	cases := []struct {
 		name string