From 39f185d881db3e24729e77e60441baf537e97047 Mon Sep 17 00:00:00 2001
From: German Lashevich <german.lashevich@gmail.com>
Date: Sun, 2 Feb 2025 16:51:21 +0100
Subject: [PATCH] fix: correctly assign apps to nested envs (#389)

---
 internal/myks/globe.go                        |  17 ++-
 internal/myks/globe_test.go                   | 115 ++++++++++++++++++
 .../envs/bar/prod/env-data.ytt.yaml           |   2 +
 .../envs/bar/stage/env-data.ytt.yaml          |   2 +
 .../envs/foo/prod/env-data.ytt.yaml           |   2 +
 .../envs/foo/stage/env-data.ytt.yaml          |   2 +
 6 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 testData/collect-environments/envs/bar/prod/env-data.ytt.yaml
 create mode 100644 testData/collect-environments/envs/bar/stage/env-data.ytt.yaml
 create mode 100644 testData/collect-environments/envs/foo/prod/env-data.ytt.yaml
 create mode 100644 testData/collect-environments/envs/foo/stage/env-data.ytt.yaml

diff --git a/internal/myks/globe.go b/internal/myks/globe.go
index c6ccf6c5..c687cb31 100644
--- a/internal/myks/globe.go
+++ b/internal/myks/globe.go
@@ -196,6 +196,7 @@ func (g *Globe) ValidateRootDir() error {
 
 func (g *Globe) Init(asyncLevel int, envSearchPathToAppMap EnvAppMap) error {
 	envAppMap := g.collectEnvironments(g.AddBaseDirToEnvAppMap(envSearchPathToAppMap))
+	log.Debug().Interface("envAppMap", envAppMap).Msg(g.Msg("Environments collected from search paths"))
 
 	return process(asyncLevel, maps.Keys(envAppMap), func(item interface{}) error {
 		envPath, ok := item.(string)
@@ -391,9 +392,23 @@ func (g *Globe) collectEnvironments(envSearchPathToAppMap EnvAppMap) EnvAppMap {
 
 	for searchPath, appNames := range envSearchPathToAppMap {
 		for _, envPath := range g.collectEnvironmentsInPath(searchPath) {
-			envAppMap[envPath] = appNames
+			// If appNames is nil or empty, it means all applications in the environment should be processed
+			if len(appNames) == 0 {
+				envAppMap[envPath] = nil
+				continue
+			}
+			// If the environment is already in the map, append the appNames to the existing list,
+			// but only if its apps are not nil
+			if apps, ok := envAppMap[envPath]; !ok {
+				envAppMap[envPath] = appNames
+			} else if apps != nil {
+				envAppMap[envPath] = append(apps, appNames...)
+			}
 		}
 	}
+	for env, apps := range envAppMap {
+		envAppMap[env] = unique(apps)
+	}
 
 	log.Debug().Interface("envToAppMap", envAppMap).Msg(g.Msg("Collected environments"))
 	return envAppMap
diff --git a/internal/myks/globe_test.go b/internal/myks/globe_test.go
index cdc4f2e4..add6c01d 100644
--- a/internal/myks/globe_test.go
+++ b/internal/myks/globe_test.go
@@ -1,6 +1,7 @@
 package myks
 
 import (
+	"sort"
 	"testing"
 )
 
@@ -26,3 +27,117 @@ func Test_AddBaseDirToEnvPath(t *testing.T) {
 		}
 	}
 }
+
+func Test_collectEnvironments(t *testing.T) {
+	defer chdir(t, "../../testData/collect-environments")()
+
+	g := NewWithDefaults()
+	g.environments = make(map[string]*Environment)
+	tests := []struct {
+		name string
+		in   EnvAppMap
+		out  EnvAppMap
+	}{
+		{
+			"empty map",
+			EnvAppMap{},
+			EnvAppMap{
+				"envs/bar/prod":  []string{},
+				"envs/bar/stage": []string{},
+				"envs/foo/prod":  []string{},
+				"envs/foo/stage": []string{},
+			},
+		},
+		{
+			"one env",
+			EnvAppMap{
+				"envs/bar/prod": []string{"app1", "app2"},
+			},
+			EnvAppMap{
+				"envs/bar/prod": []string{"app1", "app2"},
+			},
+		},
+		{
+			"multiple envs",
+			EnvAppMap{
+				"envs/bar/prod":  []string{"app1", "app2"},
+				"envs/bar/stage": []string{"app1", "app3"},
+				"envs/foo/prod":  []string{"app1"},
+			},
+			EnvAppMap{
+				"envs/bar/prod":  []string{"app1", "app2"},
+				"envs/bar/stage": []string{"app1", "app3"},
+				"envs/foo/prod":  []string{"app1"},
+			},
+		},
+		{
+			"nested envs",
+			EnvAppMap{
+				"envs":           []string{"app1", "envsApp"},
+				"envs/bar":       []string{"app2", "barApp"},
+				"envs/bar/stage": []string{"app3", "stageApp"},
+			},
+			EnvAppMap{
+				"envs/bar/prod":  []string{"app1", "envsApp", "app2", "barApp"},
+				"envs/bar/stage": []string{"app1", "envsApp", "app2", "barApp", "app3", "stageApp"},
+				"envs/foo/prod":  []string{"app1", "envsApp"},
+				"envs/foo/stage": []string{"app1", "envsApp"},
+			},
+		},
+		{
+			"deduplication",
+			EnvAppMap{
+				"envs/bar": []string{"app1", "app1", "app2"},
+			},
+			EnvAppMap{
+				"envs/bar/prod":  []string{"app1", "app2"},
+				"envs/bar/stage": []string{"app1", "app2"},
+			},
+		},
+		{
+			"empty list prioritised",
+			EnvAppMap{
+				"envs/bar":      nil,
+				"envs/bar/prod": []string{"app1", "app2"},
+				"envs/foo":      []string{},
+				"envs/foo/prod": []string{"app3"},
+			},
+			EnvAppMap{
+				"envs/bar/prod":  []string{},
+				"envs/bar/stage": []string{},
+				"envs/foo/prod":  []string{},
+				"envs/foo/stage": []string{},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		out := g.collectEnvironments(tt.in)
+		if !compareEnvAppMap(out, tt.out) {
+			t.Errorf("%s:\n  got  %v\n  want %v", tt.name, out, tt.out)
+		}
+	}
+}
+
+func compareEnvAppMap(left, right EnvAppMap) bool {
+	if len(left) != len(right) {
+		return false
+	}
+	for leftEnv, leftApps := range left {
+		if rightApps, ok := right[leftEnv]; !ok {
+			return false
+		} else {
+			if len(leftApps) != len(rightApps) {
+				return false
+			}
+			sort.Strings(leftApps)
+			sort.Strings(rightApps)
+			for i, leftApp := range leftApps {
+				if leftApp != rightApps[i] {
+					return false
+				}
+			}
+		}
+	}
+	return true
+}
diff --git a/testData/collect-environments/envs/bar/prod/env-data.ytt.yaml b/testData/collect-environments/envs/bar/prod/env-data.ytt.yaml
new file mode 100644
index 00000000..f2842f04
--- /dev/null
+++ b/testData/collect-environments/envs/bar/prod/env-data.ytt.yaml
@@ -0,0 +1,2 @@
+environment:
+  id: bar-prod
diff --git a/testData/collect-environments/envs/bar/stage/env-data.ytt.yaml b/testData/collect-environments/envs/bar/stage/env-data.ytt.yaml
new file mode 100644
index 00000000..4ad9099e
--- /dev/null
+++ b/testData/collect-environments/envs/bar/stage/env-data.ytt.yaml
@@ -0,0 +1,2 @@
+environment:
+  id: bar-stage
diff --git a/testData/collect-environments/envs/foo/prod/env-data.ytt.yaml b/testData/collect-environments/envs/foo/prod/env-data.ytt.yaml
new file mode 100644
index 00000000..7d1b9ff8
--- /dev/null
+++ b/testData/collect-environments/envs/foo/prod/env-data.ytt.yaml
@@ -0,0 +1,2 @@
+environment:
+  id: foo-prod
diff --git a/testData/collect-environments/envs/foo/stage/env-data.ytt.yaml b/testData/collect-environments/envs/foo/stage/env-data.ytt.yaml
new file mode 100644
index 00000000..d484fc1c
--- /dev/null
+++ b/testData/collect-environments/envs/foo/stage/env-data.ytt.yaml
@@ -0,0 +1,2 @@
+environment:
+  id: foo-stage