flannel/subnet/subnet_test.go

package subnet

import (
	"fmt"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/coreos-inc/kolach/Godeps/_workspace/src/github.com/coreos/go-etcd/etcd"

	"github.com/coreos-inc/kolach/pkg"
)

type mockSubnetRegistry struct {
	subnets *etcd.Node
	ch      chan string
	index   uint64
}

func newMockSubnetRegistry(ch chan string) *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},
		&etcd.Node{Key: "10.3.4.0-24", Value: `{ "PublicIP": "1.1.1.1" }`, ModifiedIndex: 12},
		&etcd.Node{Key: "10.3.5.0-24", Value: `{ "PublicIP": "1.1.1.1" }`, ModifiedIndex: 13},
	}

	return &mockSubnetRegistry{
		subnets: &etcd.Node{
			Nodes: subnodes,
		},
		ch:    ch,
		index: 14,
	}
}

func (msr *mockSubnetRegistry) getConfig() (*etcd.Response, error) {
	return &etcd.Response{
		EtcdIndex: msr.index,
		Node: &etcd.Node{
			Value: `{ "Network": "10.3.0.0/16", "FirstIP": "10.3.1.0", "LastIP": "10.3.5.0" }`,
		},
	}, nil
}

func (msr *mockSubnetRegistry) getSubnets() (*etcd.Response, error) {
	return &etcd.Response{
		Node:      msr.subnets,
		EtcdIndex: msr.index,
	}, nil
}

func (msr *mockSubnetRegistry) createSubnet(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,
	}

	msr.subnets.Nodes = append(msr.subnets.Nodes, node)
	return &etcd.Response{
		Node:      node,
		EtcdIndex: msr.index,
	}, nil
}

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,
	}

	return &etcd.Response{
		Node:      node,
		EtcdIndex: msr.index,
	}, nil
}

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,
			}
			msr.subnets.Nodes = append(msr.subnets.Nodes, &n)
			return &etcd.Response{Node: &n}, nil
		}
	}
}

func (msr *mockSubnetRegistry) hasSubnet(sn string) bool {
	for _, n := range msr.subnets.Nodes {
		if n.Key == sn {
			return true
		}
	}
	return false
}

func netIPNetToString(n *net.IPNet) string {
	return strings.Replace(n.String(), "/", "-", 1)
}

func TestAcquireLease(t *testing.T) {
	msr := newMockSubnetRegistry(nil)
	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)
	}

	if sn.String() != "10.3.3.0/24" {
		t.Fatal("Subnet mismatch: expected 10.3.3.0/24, got: ", sn)
	}

	// Acquire again, should reuse
	if sn, err = sm.AcquireLease(ip, data); err != nil {
		t.Fatal("AcquireLease failed: ", err)
	}

	if sn.String() != "10.3.3.0/24" {
		t.Fatal("Subnet mismatch: expected 10.3.3.0/24, got: ", sn)
	}
}

func TestWatchLeases(t *testing.T) {
	msr := newMockSubnetRegistry(make(chan string))
	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" }`

	_, err = sm.AcquireLease(ip, data)
	if err != nil {
		t.Fatalf("RegisterSubnet failed: %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
		}
	}

	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()
}