From 60443824dd511279fd438b6cd1ce4265e8af1d21 Mon Sep 17 00:00:00 2001
From: Evgeny Abramovich <evg.abramovitch@gmail.com>
Date: Sat, 19 Nov 2022 13:41:58 -0400
Subject: [PATCH] Updated new url mapping loading

---
 internal/config/helpers.go      | 31 ++++++++++++++
 internal/config/helpers_test.go | 71 +++++++++++++++++++++++++++++++++
 main.go                         | 11 +++--
 3 files changed, 110 insertions(+), 3 deletions(-)
 create mode 100644 internal/config/helpers.go
 create mode 100644 internal/config/helpers_test.go

diff --git a/internal/config/helpers.go b/internal/config/helpers.go
new file mode 100644
index 00000000..6fed77e2
--- /dev/null
+++ b/internal/config/helpers.go
@@ -0,0 +1,31 @@
+package config
+
+import (
+	"errors"
+
+	"github.com/spf13/viper"
+)
+
+var (
+	ErrNoToPair   = errors.New("`to` values are not set for every `from`")
+	ErrNoFromPair = errors.New("`from` values are not set for every `to`")
+)
+
+func ReadURLMapping(config *viper.Viper) (map[string]string, error) {
+	urlMappings := map[string]string{}
+	from, to := config.GetStringSlice("from"), config.GetStringSlice("to") // nolint: varnamelen
+
+	if len(from) > len(to) {
+		return nil, ErrNoToPair
+	}
+
+	if len(to) > len(from) {
+		return nil, ErrNoFromPair
+	}
+
+	for index, key := range from {
+		urlMappings[key] = to[index]
+	}
+
+	return urlMappings, nil
+}
diff --git a/internal/config/helpers_test.go b/internal/config/helpers_test.go
new file mode 100644
index 00000000..2e86c48d
--- /dev/null
+++ b/internal/config/helpers_test.go
@@ -0,0 +1,71 @@
+package config_test
+
+import (
+	"testing"
+
+	"github.com/evg4b/uncors/internal/config"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReadURLMapping(t *testing.T) {
+	t.Run("correctly map pairs", func(t *testing.T) {
+		viperConfig := viper.New()
+		viperConfig.Set("from", []string{"host1", "host2", "host3"})
+		viperConfig.Set("to", []string{"target-host1", "target-host2", "target-host3"})
+		actual, err := config.ReadURLMapping(viperConfig)
+
+		assert.NoError(t, err)
+		assert.EqualValues(t, map[string]string{
+			"host1": "target-host1",
+			"host2": "target-host2",
+			"host3": "target-host3",
+		}, actual)
+	})
+
+	t.Run("incorrect pairs", func(t *testing.T) {
+		tests := []struct {
+			name        string
+			from        []string
+			to          []string
+			expectedErr string
+		}{
+			{
+				name:        "from is not set",
+				from:        []string{"host1"},
+				to:          []string{},
+				expectedErr: "`to` values are not set for every `from`",
+			},
+			{
+				name:        "to is not set",
+				from:        []string{},
+				to:          []string{"target-host1"},
+				expectedErr: "`from` values are not set for every `to`",
+			},
+			{
+				name:        "count of from values greath then count of to",
+				from:        []string{"host1", "host2"},
+				to:          []string{"target-host1"},
+				expectedErr: "`to` values are not set for every `from`",
+			},
+			{
+				name:        "count of to values greath then count of from",
+				from:        []string{"host1"},
+				to:          []string{"target-host1", "target-host2"},
+				expectedErr: "`from` values are not set for every `to`",
+			},
+		}
+		for _, testCase := range tests {
+			t.Run(testCase.name, func(t *testing.T) {
+				viperConfig := viper.New()
+				viperConfig.Set("from", testCase.from)
+				viperConfig.Set("to", testCase.to)
+
+				actual, err := config.ReadURLMapping(viperConfig)
+
+				assert.Nil(t, actual)
+				assert.EqualError(t, err, testCase.expectedErr)
+			})
+		}
+	})
+}
diff --git a/main.go b/main.go
index da42c524..31e2f120 100644
--- a/main.go
+++ b/main.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 
+	"github.com/evg4b/uncors/internal/config"
 	"github.com/evg4b/uncors/internal/infrastructure"
 	"github.com/evg4b/uncors/internal/log"
 	"github.com/evg4b/uncors/internal/mock"
@@ -29,8 +30,8 @@ const (
 )
 
 func main() {
-	pflag.String("target", "https://github.com", "Target host with protocol for to the resource to be proxy")
-	pflag.String("source", "localhost", "Local host with protocol for to the resource from which proxying will take place") // nolint: lll
+	pflag.StringSlice("to", []string{}, "Target host with protocol for to the resource to be proxy")
+	pflag.StringSlice("from", []string{}, "Local host with protocol for to the resource from which proxying will take place") // nolint: lll
 	pflag.Int("http-port", defaultHTTPPort, "Local HTTP listening port")
 	pflag.Int("https-port", defaultHTTPSPort, "Local HTTPS listening port")
 	pflag.String("cert-file", "", "Path to HTTPS certificate file")
@@ -80,8 +81,12 @@ func main() {
 
 	mock.MakeMockedRoutes(router, ui.MockLogger, mocksDefs)
 
+	urlMappings, err := config.ReadURLMapping(viper.GetViper())
+	if err != nil {
+		panic(err)
+	}
 	mappings, err := urlreplacer.NormaliseMappings(
-		map[string]string{viper.GetString("source"): viper.GetString("target")},
+		urlMappings,
 		httpPort,
 		httpsPort,
 		len(certFile) > 0 && len(keyFile) > 0,