From f49bdcbf33a781fc142fa045fa38d803eeb5a621 Mon Sep 17 00:00:00 2001
From: Solomon Hykes <solomon@dotcloud.com>
Date: Fri, 5 Apr 2013 13:03:24 -0700
Subject: [PATCH] Add tests of tcp port allocator Upstream-commit:
 febaeebfb8848265267213b2f6a6fc3a40ad90f1 Component: engine

---
 components/engine/network.go      |  2 ++
 components/engine/network_test.go | 36 +++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/components/engine/network.go b/components/engine/network.go
index 5b4f312263..982559868f 100644
--- a/components/engine/network.go
+++ b/components/engine/network.go
@@ -176,6 +176,7 @@ func (alloc *PortAllocator) runFountain() {
 
 // FIXME: Release can no longer fail, change its prototype to reflect that.
 func (alloc *PortAllocator) Release(port int) error {
+	Debugf("Releasing %d", port)
 	alloc.lock.Lock()
 	delete(alloc.inUse, port)
 	alloc.lock.Unlock()
@@ -183,6 +184,7 @@ func (alloc *PortAllocator) Release(port int) error {
 }
 
 func (alloc *PortAllocator) Acquire(port int) (int, error) {
+	Debugf("Acquiring %d", port)
 	if port == 0 {
 		// Allocate a port from the fountain
 		for port := range alloc.fountain {
diff --git a/components/engine/network_test.go b/components/engine/network_test.go
index a9d3cac454..8f31d27a98 100644
--- a/components/engine/network_test.go
+++ b/components/engine/network_test.go
@@ -18,6 +18,42 @@ func TestIptables(t *testing.T) {
 	}
 }
 
+func TestPortAllocation(t *testing.T) {
+	allocator, err := newPortAllocator()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port, err := allocator.Acquire(80); err != nil {
+		t.Fatal(err)
+	} else if port != 80 {
+		t.Fatalf("Acquire(80) should return 80, not %d", port)
+	}
+	port, err := allocator.Acquire(0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if port <= 0 {
+		t.Fatalf("Acquire(0) should return a non-zero port")
+	}
+	if _, err := allocator.Acquire(port); err == nil {
+		t.Fatalf("Acquiring a port already in use should return an error")
+	}
+	if newPort, err := allocator.Acquire(0); err != nil {
+		t.Fatal(err)
+	} else if newPort == port {
+		t.Fatalf("Acquire(0) allocated the same port twice: %d", port)
+	}
+	if _, err := allocator.Acquire(80); err == nil {
+		t.Fatalf("Acquiring a port already in use should return an error")
+	}
+	if err := allocator.Release(80); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := allocator.Acquire(80); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestNetworkRange(t *testing.T) {
 	// Simple class C test
 	_, network, _ := net.ParseCIDR("192.168.0.1/24")