From de093d80ea1a2210a77e715e2886b5380adc18c3 Mon Sep 17 00:00:00 2001
From: Evgeniy Kulikov <im@kulikov.im>
Date: Sun, 10 Nov 2019 00:33:00 +0300
Subject: [PATCH] Add possibility to fetch nats.servers from environments

---
 nats.go      | 27 +++++++++++++++++++++++++--
 nats_test.go | 20 ++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/nats.go b/nats.go
index b622b28..2cf58f2 100644
--- a/nats.go
+++ b/nats.go
@@ -1,6 +1,8 @@
 package nats
 
 import (
+	"strconv"
+
 	"github.com/im-kulikov/helium/module"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/stan.go"
@@ -61,6 +63,27 @@ var (
 	}
 )
 
+func fetchAddresses(key string, v *viper.Viper) []string {
+	var (
+		addresses []string
+	)
+
+	for i := 0; ; i++ {
+		addr := v.GetString(key + "_" + strconv.Itoa(i))
+		if addr == "" {
+			break
+		}
+
+		addresses = append(addresses, addr)
+	}
+
+	if len(addresses) == 0 {
+		addresses = v.GetStringSlice(key)
+	}
+
+	return addresses
+}
+
 // NewDefaultConfig default settings for connection
 func NewDefaultConfig(v *viper.Viper) (*Config, error) {
 	if !v.IsSet("nats") {
@@ -68,8 +91,8 @@ func NewDefaultConfig(v *viper.Viper) (*Config, error) {
 	}
 
 	var servers []string
-	if v.IsSet("nats.servers") {
-		servers = v.GetStringSlice("nats.servers")
+	if addresses := fetchAddresses("nats.servers", v); len(addresses) > 0 {
+		servers = addresses
 	}
 
 	return &Config{
diff --git a/nats_test.go b/nats_test.go
index e0cc097..7206945 100644
--- a/nats_test.go
+++ b/nats_test.go
@@ -55,6 +55,26 @@ func TestNewDefaultConfig(t *testing.T) {
 		require.Equal(t, c.Url, url)
 	})
 
+	t.Run("should fetch nats.servers_0", func(t *testing.T) {
+		v := viper.New()
+		v.SetDefault("nats.servers_0", nats.DefaultURL)
+
+		c, err := NewDefaultConfig(v)
+		require.NoError(t, err)
+		require.Len(t, c.Servers, 1)
+		require.Equal(t, c.Servers[0], nats.DefaultURL)
+	})
+
+	t.Run("should fetch nats.servers", func(t *testing.T) {
+		v := viper.New()
+		v.SetDefault("nats.servers", []string{nats.DefaultURL})
+
+		c, err := NewDefaultConfig(v)
+		require.NoError(t, err)
+		require.Len(t, c.Servers, 1)
+		require.Equal(t, c.Servers[0], nats.DefaultURL)
+	})
+
 	t.Run("should fail for empty config", func(t *testing.T) {
 		c, err := NewConnection(nil)
 		require.Nil(t, c)