From 439b7e8b9880b116bd1ecd7df803ef1bb2c4afd1 Mon Sep 17 00:00:00 2001
From: Jonathan Boulle <jonathanboulle@gmail.com>
Date: Tue, 24 Jun 2014 18:49:49 -0700
Subject: [PATCH] initialize/env: fall back to COREOS_*_IPV4 env variables

---
 initialize/env.go      | 10 ++++++++
 initialize/env_test.go | 52 ++++++++++++++++++++++++++++++------------
 2 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/initialize/env.go b/initialize/env.go
index b2816e2..715c07d 100644
--- a/initialize/env.go
+++ b/initialize/env.go
@@ -1,6 +1,7 @@
 package initialize
 
 import (
+	"os"
 	"path"
 	"strings"
 )
@@ -18,6 +19,15 @@ type Environment struct {
 
 // TODO(jonboulle): this is getting unwieldy, should be able to simplify the interface somehow
 func NewEnvironment(root, configRoot, workspace, netconfType, sshKeyName string, substitutions map[string]string) *Environment {
+	// If certain values are not in the supplied substitution, fall back to retrieving them from the environment
+	for k, v := range map[string]string{
+		"$public_ipv4":  os.Getenv("COREOS_PUBLIC_IPV4"),
+		"$private_ipv4": os.Getenv("COREOS_PRIVATE_IPV4"),
+	} {
+		if _, ok := substitutions[k]; !ok {
+			substitutions[k] = v
+		}
+	}
 	return &Environment{root, configRoot, workspace, netconfType, sshKeyName, substitutions}
 }
 
diff --git a/initialize/env_test.go b/initialize/env_test.go
index 3ff28de..8b572ac 100644
--- a/initialize/env_test.go
+++ b/initialize/env_test.go
@@ -1,26 +1,48 @@
 package initialize
 
-import "testing"
+import (
+	"os"
+	"testing"
+)
 
 func TestEnvironmentApply(t *testing.T) {
-	subs := map[string]string{
-		"$public_ipv4":  "192.0.2.3",
-		"$private_ipv4": "192.0.2.203",
-	}
-	env := NewEnvironment("./", "./", "./", "", "", subs)
-	input := `[Service]
+	os.Setenv("COREOS_PUBLIC_IPV4", "1.2.3.4")
+	os.Setenv("COREOS_PRIVATE_IPV4", "5.6.7.8")
+	for _, tt := range []struct {
+		subs  map[string]string
+		input string
+		out   string
+	}{
+		{
+			map[string]string{
+				"$public_ipv4":  "192.0.2.3",
+				"$private_ipv4": "192.0.2.203",
+			},
+			`[Service]
 ExecStart=/usr/bin/echo "$public_ipv4"
 ExecStop=/usr/bin/echo $private_ipv4
-ExecStop=/usr/bin/echo $unknown
-`
-	expected := `[Service]
+ExecStop=/usr/bin/echo $unknown`,
+			`[Service]
 ExecStart=/usr/bin/echo "192.0.2.3"
 ExecStop=/usr/bin/echo 192.0.2.203
-ExecStop=/usr/bin/echo $unknown
-`
+ExecStop=/usr/bin/echo $unknown`,
+		},
+		{
+			map[string]string{"$private_ipv4": "127.0.0.1"},
+			"$private_ipv4\n$public_ipv4",
+			"127.0.0.1\n1.2.3.4",
+		},
+		{
+			map[string]string{"foo": "bar"},
+			"$private_ipv4\n$public_ipv4",
+			"5.6.7.8\n1.2.3.4",
+		},
+	} {
 
-	output := env.Apply(input)
-	if output != expected {
-		t.Fatalf("Environment incorrectly applied.\nOutput:\n%s\nExpected:\n%s", output, expected)
+		env := NewEnvironment("./", "./", "./", "", "", tt.subs)
+		got := env.Apply(tt.input)
+		if got != tt.out {
+			t.Fatalf("Environment incorrectly applied.\ngot:\n%s\nwant:\n%s", got, tt.out)
+		}
 	}
 }