From ba9f1016608a40f5e49c99bf36b7e2beaba2d70b Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Sun, 15 Nov 2015 01:13:42 -0800
Subject: [PATCH] Fixed the test related to setting env variables in tasks

---
 client/driver/driver.go      | 1 +
 client/driver/driver_test.go | 8 ++++++--
 nomad/structs/structs.go     | 9 +++++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/client/driver/driver.go b/client/driver/driver.go
index 080872690..7fc76863f 100644
--- a/client/driver/driver.go
+++ b/client/driver/driver.go
@@ -133,6 +133,7 @@ func TaskEnvironmentVariables(ctx *ExecContext, task *structs.Task) environment.
 		if len(task.Resources.Networks) > 0 {
 			network := task.Resources.Networks[0]
 			env.SetTaskIp(network.IP)
+			env.SetPorts(network.MapLabelToValues())
 		}
 	}
 
diff --git a/client/driver/driver_test.go b/client/driver/driver_test.go
index 76877d133..7065153a1 100644
--- a/client/driver/driver_test.go
+++ b/client/driver/driver_test.go
@@ -76,8 +76,12 @@ func TestDriver_TaskEnvironmentVariables(t *testing.T) {
 		"NOMAD_CPU_LIMIT":       "1000",
 		"NOMAD_MEMORY_LIMIT":    "500",
 		"NOMAD_IP":              "1.2.3.4",
-		"NOMAD_PORT_admin":      "8080",
-		"NOMAD_PORT_5000":       "12345",
+		"NOMAD_PORT_one":        "80",
+		"NOMAD_PORT_two":        "443",
+		"NOMAD_PORT_three":      "8080",
+		"NOMAD_PORT_four":       "12345",
+		"NOMAD_PORT_admin":      "8081",
+		"NOMAD_PORT_web":        "8086",
 		"NOMAD_META_CHOCOLATE":  "cake",
 		"NOMAD_META_STRAWBERRY": "icecream",
 		"HELLO":                 "world",
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 1a8f3b134..b9a326e0f 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -661,6 +661,15 @@ func (n *NetworkResource) GoString() string {
 	return fmt.Sprintf("*%#v", *n)
 }
 
+func (n *NetworkResource) MapLabelToValues() map[string]int {
+	labelValues := make(map[string]int)
+	ports := append(n.ReservedPorts, n.DynamicPorts...)
+	for _, port := range ports {
+		labelValues[port.Label] = port.Value
+	}
+	return labelValues
+}
+
 const (
 	// JobTypeNomad is reserved for internal system tasks and is
 	// always handled by the CoreScheduler.