From a171aa4828f6fd5b8dae09dc6903074eb5bc3be1 Mon Sep 17 00:00:00 2001
From: Chris Marslender <chrismarslender@gmail.com>
Date: Thu, 1 Aug 2024 14:22:44 -0500
Subject: [PATCH] Add wonkydict to account for default empty values of either
 [] or {}

---
 pkg/config/config.go       | 18 +++++++--------
 pkg/types/wonkyset.go      | 25 +++++++++++++++++++++
 pkg/types/wonkyset_test.go | 45 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 pkg/types/wonkyset.go
 create mode 100644 pkg/types/wonkyset_test.go

diff --git a/pkg/config/config.go b/pkg/config/config.go
index aec78db..f1cde2b 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -194,15 +194,15 @@ type PoolConfig struct {
 
 // FarmerConfig farmer configuration section
 type FarmerConfig struct {
-	FullNodePeers      []Peer            `yaml:"full_node_peers"`
-	PoolPublicKeys     map[string]string `yaml:"pool_public_keys"`
-	XCHTargetAddress   string            `yaml:"xch_target_address,omitempty"`
-	StartRPCServer     bool              `yaml:"start_rpc_server"`
-	EnableProfiler     bool              `yaml:"enable_profiler"`
-	PoolShareThreshold uint32            `yaml:"pool_share_threshold"`
-	Logging            LoggingConfig     `yaml:"logging"`
-	NetworkOverrides   NetworkOverrides  `yaml:"network_overrides"`
-	SelectedNetwork    string            `yaml:"selected_network"`
+	FullNodePeers      []Peer           `yaml:"full_node_peers"`
+	PoolPublicKeys     types.WonkySet   `yaml:"pool_public_keys"`
+	XCHTargetAddress   string           `yaml:"xch_target_address,omitempty"`
+	StartRPCServer     bool             `yaml:"start_rpc_server"`
+	EnableProfiler     bool             `yaml:"enable_profiler"`
+	PoolShareThreshold uint32           `yaml:"pool_share_threshold"`
+	Logging            LoggingConfig    `yaml:"logging"`
+	NetworkOverrides   NetworkOverrides `yaml:"network_overrides"`
+	SelectedNetwork    string           `yaml:"selected_network"`
 	PortConfig         `yaml:",inline"`
 	SSL                SSLConfig `yaml:"ssl"`
 }
diff --git a/pkg/types/wonkyset.go b/pkg/types/wonkyset.go
new file mode 100644
index 0000000..9102842
--- /dev/null
+++ b/pkg/types/wonkyset.go
@@ -0,0 +1,25 @@
+package types
+
+// WonkySet is just an alias for map[string]string that will unmarshal correctly from the empty forms {} and []
+// In chia-blockchain, the default for some sets was incorrectly [] in the initial config
+// so this ensures compatability with both ways the empty values could show up
+type WonkySet map[string]string
+
+// UnmarshalYAML implements the yaml.Unmarshaler interface.
+func (ws *WonkySet) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	// Attempt to unmarshal into a slice of strings.
+	var slice []string
+	if err := unmarshal(&slice); err == nil {
+		*ws = make(map[string]string)
+		return nil
+	}
+
+	// Attempt to unmarshal into a map.
+	var m map[string]string
+	if err := unmarshal(&m); err != nil {
+		return err
+	}
+
+	*ws = m
+	return nil
+}
diff --git a/pkg/types/wonkyset_test.go b/pkg/types/wonkyset_test.go
new file mode 100644
index 0000000..e0b2798
--- /dev/null
+++ b/pkg/types/wonkyset_test.go
@@ -0,0 +1,45 @@
+package types_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/yaml.v3"
+
+	"github.com/chia-network/go-chia-libs/pkg/config"
+)
+
+// TestWonkySet_UnmarshalYAML ensures that this will unmarshal both empty list [] and dict {} into empty map[string]string
+// And also ensures an actal !!set as it would show up in the yaml unmarshals correctly
+func TestWonkySet_UnmarshalYAML(t *testing.T) {
+	var yamlWithList = []byte(`
+farmer:
+  pool_public_keys: []
+`)
+	var yamlWithDict = []byte(`
+farmer:
+  pool_public_keys: {}
+`)
+
+	var yamlWithData = []byte(`
+farmer:
+  pool_public_keys: !!set
+    abc123: null
+    456xyz: null
+`)
+
+	cfg := &config.ChiaConfig{}
+	err := yaml.Unmarshal(yamlWithList, cfg)
+	assert.NoError(t, err)
+	assert.Len(t, cfg.Farmer.PoolPublicKeys, 0)
+
+	cfg = &config.ChiaConfig{}
+	err = yaml.Unmarshal(yamlWithDict, cfg)
+	assert.NoError(t, err)
+	assert.Len(t, cfg.Farmer.PoolPublicKeys, 0)
+
+	cfg = &config.ChiaConfig{}
+	err = yaml.Unmarshal(yamlWithData, cfg)
+	assert.NoError(t, err)
+	assert.Len(t, cfg.Farmer.PoolPublicKeys, 2)
+}