Merge pull request #8 from eyakubovich/master

added subnet expiration and renewal tests

subnet/subnet.go

@@ -85,7 +85,7 @@ func (sm *SubnetManager) AcquireLease(ip pkg.IP4, data string) (pkg.IP4Net, erro
 					}
 
 					sm.myLease.Network = l.Network
-					sm.leaseExp = *(resp.Node.Expiration)
+					sm.leaseExp = *resp.Node.Expiration
 					return l.Network, nil
 				}
 			}
@@ -101,7 +101,7 @@ func (sm *SubnetManager) AcquireLease(ip pkg.IP4, data string) (pkg.IP4Net, erro
 		switch {
 		case err == nil:
 			sm.myLease.Network = sn
-			sm.leaseExp = *(resp.Node.Expiration)
+			sm.leaseExp = *resp.Node.Expiration
 			return sn, nil
 
 		// if etcd returned Key Already Exists, try again.
@@ -118,7 +118,7 @@ func (sm *SubnetManager) AcquireLease(ip pkg.IP4, data string) (pkg.IP4Net, erro
 
 func (sm *SubnetManager) UpdateSubnet(data string) error {
 	resp, err := sm.registry.updateSubnet(sm.myLease.Network.StringSep(".", "-"), data, subnetTTL)
-	sm.leaseExp = *(resp.Node.Expiration)
+	sm.leaseExp = *resp.Node.Expiration
 	return err
 }
 

subnet/subnet_test.go

@@ -2,8 +2,6 @@ package subnet
 
 import (
 	"fmt"
-	"net"
-	"strings"
 	"testing"
 	"time"
 
@@ -14,11 +12,13 @@ import (
 
 type mockSubnetRegistry struct {
 	subnets *etcd.Node
-	ch      chan string
+	addCh   chan string
+	delCh   chan string
 	index   uint64
+	ttl     uint64
 }
 
-func newMockSubnetRegistry(ch chan string) *mockSubnetRegistry {
+func newMockSubnetRegistry(ttlOverride uint64) *mockSubnetRegistry {
 	subnodes := []*etcd.Node{
 		&etcd.Node{Key: "10.3.1.0-24", Value: `{ "PublicIP": "1.1.1.1" }`, ModifiedIndex: 10},
 		&etcd.Node{Key: "10.3.2.0-24", Value: `{ "PublicIP": "1.1.1.1" }`, ModifiedIndex: 11},
@@ -30,8 +30,10 @@ func newMockSubnetRegistry(ch chan string) *mockSubnetRegistry {
 		subnets: &etcd.Node{
 			Nodes: subnodes,
 		},
-		ch:    ch,
+		addCh: make(chan string),
+		delCh: make(chan string),
 		index: 14,
+		ttl:   ttlOverride,
 	}
 }
 
@@ -54,6 +56,10 @@ func (msr *mockSubnetRegistry) getSubnets() (*etcd.Response, error) {
 func (msr *mockSubnetRegistry) createSubnet(sn, data string, ttl uint64) (*etcd.Response, error) {
 	msr.index += 1
 
+	if msr.ttl > 0 {
+		ttl = msr.ttl
+	}
+
 	// add squared durations :)
 	exp := time.Now().Add(time.Duration(ttl) * time.Second)
 
@@ -65,6 +71,7 @@ func (msr *mockSubnetRegistry) createSubnet(sn, data string, ttl uint64) (*etcd.
 	}
 
 	msr.subnets.Nodes = append(msr.subnets.Nodes, node)
+
 	return &etcd.Response{
 		Node:      node,
 		EtcdIndex: msr.index,
@@ -72,38 +79,59 @@ func (msr *mockSubnetRegistry) createSubnet(sn, data string, ttl uint64) (*etcd.
 }
 
 func (msr *mockSubnetRegistry) updateSubnet(sn, data string, ttl uint64) (*etcd.Response, error) {
+
 	msr.index += 1
 
 	// add squared durations :)
 	exp := time.Now().Add(time.Duration(ttl) * time.Second)
 
-	node := &etcd.Node{
-		Key:           sn,
-		Value:         data,
-		ModifiedIndex: msr.index,
-		Expiration:    &exp,
+	for _, n := range msr.subnets.Nodes {
+		if n.Key == sn {
+			n.Value = data
+			n.ModifiedIndex = msr.index
+			n.Expiration = &exp
+
+			return &etcd.Response{
+				Node:      n,
+				EtcdIndex: msr.index,
+			}, nil
+		}
 	}
 
-	return &etcd.Response{
-		Node:      node,
-		EtcdIndex: msr.index,
-	}, nil
+	return nil, fmt.Errorf("Subnet not found")
+
 }
 
 func (msr *mockSubnetRegistry) watchSubnets(since uint64, stop chan bool) (*etcd.Response, error) {
-	for {
-		var sn string
-		select {
-		case <-stop:
-			return nil, nil
-		case sn = <-msr.ch:
-			n := etcd.Node{
-				Key:           sn,
-				ModifiedIndex: msr.index,
+	var sn string
+
+	select {
+	case <-stop:
+		return nil, nil
+
+	case sn = <-msr.addCh:
+		n := etcd.Node{
+			Key:           sn,
+			ModifiedIndex: msr.index,
+		}
+		msr.subnets.Nodes = append(msr.subnets.Nodes, &n)
+		return &etcd.Response{
+			Action: "add",
+			Node:   &n,
+		}, nil
+
+	case sn = <-msr.delCh:
+		for i, n := range msr.subnets.Nodes {
+			if n.Key == sn {
+				msr.subnets.Nodes[i] = msr.subnets.Nodes[len(msr.subnets.Nodes)-1]
+				msr.subnets.Nodes = msr.subnets.Nodes[:len(msr.subnets.Nodes)-2]
+				return &etcd.Response{
+					Action: "expire",
+					Node:    n,
+				}, nil
 			}
-			msr.subnets.Nodes = append(msr.subnets.Nodes, &n)
-			return &etcd.Response{Node: &n}, nil
 		}
+		return nil, fmt.Errorf("Subnet (%s) to delete was not found: ", sn)
 	}
 }
 
@@ -116,12 +144,8 @@ func (msr *mockSubnetRegistry) hasSubnet(sn string) bool {
 	return false
 }
 
-func netIPNetToString(n *net.IPNet) string {
-	return strings.Replace(n.String(), "/", "-", 1)
-}
-
 func TestAcquireLease(t *testing.T) {
-	msr := newMockSubnetRegistry(nil)
+	msr := newMockSubnetRegistry(0)
 	sm, err := newSubnetManager(msr)
 	if err != nil {
 		t.Fatalf("Failed to create subnet manager: %s", err)
@@ -149,34 +173,54 @@ func TestAcquireLease(t *testing.T) {
 	}
 }
 
-func TestWatchLeases(t *testing.T) {
-	msr := newMockSubnetRegistry(make(chan string))
+func TestWatchLeaseAdded(t *testing.T) {
+	msr := newMockSubnetRegistry(0)
 	sm, err := newSubnetManager(msr)
 	if err != nil {
 		t.Fatalf("Failed to create subnet manager: %s", err)
 	}
 
-	ip, _ := pkg.ParseIP4("1.2.3.4")
-	data := `{ "PublicIP": "1.2.3.4" }`
+	events := make(chan EventBatch)
+	sm.Start(events)
+
+	expected := "10.3.3.0-24"
+	msr.addCh <- expected
+
+	evtBatch, ok := <-events
+	if !ok {
+		t.Fatalf("WatchSubnets did not publish")
+	}
+
+	if len(evtBatch) != 1 {
+		t.Fatalf("WatchSubnets produced wrong sized event batch")
+	}
+
+	evt := evtBatch[0]
+
+	if evt.Type != SubnetAdded {
+		t.Fatalf("WatchSubnets produced wrong event type")
+	}
+
+	actual := evt.Lease.Network.StringSep(".", "-")
+	if actual != expected {
+		t.Errorf("WatchSubnet produced wrong subnet: expected %s, got %s", expected, actual)
+	}
+
+	sm.Stop()
+}
 
-	_, err = sm.AcquireLease(ip, data)
+func TestWatchLeaseRemoved(t *testing.T) {
+	msr := newMockSubnetRegistry(0)
+	sm, err := newSubnetManager(msr)
 	if err != nil {
-		t.Fatalf("RegisterSubnet failed: %s", err)
+		t.Fatalf("Failed to create subnet manager: %s", err)
 	}
 
 	events := make(chan EventBatch)
 	sm.Start(events)
 
-	<-events
-
-	var expected string
-	for i := 1; i <= 9; i++ {
-		expected = fmt.Sprintf("10.3.%d.0-24", i)
-		if !msr.hasSubnet(expected) {
-			msr.ch <- expected
-			break
-		}
-	}
+	expected := "10.3.4.0-24"
+	msr.delCh <- expected
 
 	evtBatch, ok := <-events
 	if !ok {
@@ -189,7 +233,7 @@ func TestWatchLeases(t *testing.T) {
 
 	evt := evtBatch[0]
 
-	if evt.Type != SubnetAdded {
+	if evt.Type != SubnetRemoved {
 		t.Fatalf("WatchSubnets produced wrong event type")
 	}
 
@@ -200,3 +244,37 @@ func TestWatchLeases(t *testing.T) {
 
 	sm.Stop()
 }
+
+func TestRenewLease(t *testing.T) {
+	msr := newMockSubnetRegistry(1)
+	sm, err := newSubnetManager(msr)
+	if err != nil {
+		t.Fatalf("Failed to create subnet manager: %s", err)
+	}
+
+	ip, _ := pkg.ParseIP4("1.2.3.4")
+	data := `{ "PublicIP": "1.2.3.4" }`
+
+	sn, err := sm.AcquireLease(ip, data)
+	if err != nil {
+		t.Fatal("AcquireLease failed: ", err)
+	}
+
+	events := make(chan EventBatch)
+	sm.Start(events)
+
+	fmt.Println("Waiting for lease to pass original expiration")
+	time.Sleep(2*time.Second)
+
+	// check that it's still good
+	for _, n := range msr.subnets.Nodes {
+		if n.Key == sn.StringSep(".", "-") {
+			if n.Expiration.Before(time.Now()) {
+				t.Fatalf("Failed to renew lease")
+			}
+			return
+		}
+	}
+
+	t.Fatalf("Failed to find acquired lease")
+}