From de45fa6c306f914cba832e39fafe4237cb4e0f96 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Sun, 27 Aug 2023 16:40:13 +0000
Subject: [PATCH] zones: Env: minor tidy up

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/zones/env.go | 56 +++++++++++++++++++++++++++++++-----------------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/pkg/zones/env.go b/pkg/zones/env.go
index c51bf83..c70698a 100644
--- a/pkg/zones/env.go
+++ b/pkg/zones/env.go
@@ -22,11 +22,23 @@ func (m *Zones) Env(export bool) *Env {
 	}
 }
 
+// Zones returns the list of Zone IDs
+func (m *Env) Zones() []int {
+	var zones []int
+
+	m.ForEachZone(func(z *Zone) bool {
+		zones = append(zones, z.ID)
+		return false
+	})
+
+	return zones
+}
+
 // WriteTo generates environment variables for shell scripts
 func (m *Env) WriteTo(w io.Writer) (int64, error) {
 	var buf bytes.Buffer
 
-	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
+	m.writeEnvVarInts(&buf, m.Zones(), "ZONES")
 	m.ForEachZone(func(z *Zone) bool {
 		m.writeEnvZone(&buf, z)
 		return false
@@ -51,14 +63,23 @@ func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
 	}
 }
 
-func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
-	var value string
+func (m *Env) writeEnvVarInts(w io.Writer, value []int, name string, args ...any) {
+	var s string
 
-	if fn != nil {
-		value = fn(m)
+	if n := len(value); n > 0 {
+		var buf bytes.Buffer
+
+		for i, v := range value {
+			if i != 0 {
+				_, _ = fmt.Fprint(&buf, " ")
+			}
+			_, _ = fmt.Fprintf(&buf, "%v", v)
+		}
+
+		s = buf.String()
 	}
 
-	m.writeEnvVar(w, value, name, args...)
+	m.writeEnvVar(w, s, name, args...)
 }
 
 func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
@@ -79,23 +100,18 @@ func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
 	}
 }
 
-func genEnvZones(m *Env) string {
-	var s []string
-
-	m.ForEachZone(func(z *Zone) bool {
-		s = append(s, fmt.Sprintf("%v", z.ID))
-		return false
-	})
+func genEnvZoneNodes(z *Zone) string {
+	if n := z.Len(); n > 0 {
+		s := make([]string, 0, n)
 
-	return strings.Join(s, " ")
-}
+		z.ForEachMachine(func(p *Machine) bool {
+			s = append(s, p.Name)
+			return false
+		})
 
-func genEnvZoneNodes(z *Zone) string {
-	s := make([]string, 0, len(z.Machines))
-	for _, p := range z.Machines {
-		s = append(s, p.Name)
+		return strings.Join(s, " ")
 	}
-	return strings.Join(s, " ")
+	return ""
 }
 
 func getRingZeroGatewayID(z *Zone) int {