Add unit tests

controller/app.go

@@ -2,9 +2,10 @@ package controller
 
 import (
 	"fmt"
-	"github.com/golang/glog"
 	"net"
 	"strings"
+
+	"github.com/golang/glog"
 )
 
 type MonitorType int

controller/app_test.go

@@ -0,0 +1,34 @@
+package controller
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAppParsing(t *testing.T) {
+	a := assert.New(t)
+	app1, err := NewApp("app1", "1.1.1.1/32", []string{"port:tcp:123"}, []string{}, "")
+	a.Nil(err)
+	app2, err := NewApp("app1", "1.1.1.1/32", []string{"port:tcp:123"}, []string{}, "")
+	a.Nil(err)
+	app3, err := NewApp("app3", "2.2.2.2/32", []string{"exec:/bin/testme"}, []string{}, "")
+	a.Nil(err)
+
+	a.Equal("1.1.1.1/32", app1.Vip.String())
+	a.Equal(Monitor_PORT, app1.Monitors[0].Type)
+	a.Equal("123", app1.Monitors[0].Port)
+	a.Equal("tcp", app1.Monitors[0].Protocol)
+
+	a.Equal(true, app1.Equal(app2))
+
+	a.Equal(Monitor_EXEC, app3.Monitors[0].Type)
+	a.Equal("/bin/testme", app3.Monitors[0].Cmd)
+
+	// test errors
+	_, err = NewApp("app4", "4.4.4.4", []string{}, []string{}, "")
+	a.NotNil(err)
+
+	_, err = NewApp("app4", "4.4.4.4/32", []string{"port:abcd::1023"}, []string{}, "")
+	a.NotNil(err)
+}

controller/bgp.go

@@ -3,14 +3,15 @@ package controller
 import (
 	"context"
 	"fmt"
+	"net"
+	"strconv"
+	"strings"
+
 	"github.com/golang/protobuf/ptypes"
 	"github.com/golang/protobuf/ptypes/any"
 	c "github.com/mayuresh82/gocast/config"
 	api "github.com/osrg/gobgp/api"
 	gobgp "github.com/osrg/gobgp/pkg/server"
-	"net"
-	"strconv"
-	"strings"
 )
 
 type Controller struct {
@@ -22,22 +23,33 @@ type Controller struct {
 	s               *gobgp.BgpServer
 }
 
-func NewController(config *c.Config) (*Controller, error) {
+func NewController(config c.BgpConfig) (*Controller, error) {
 	c := &Controller{}
 	var gw net.IP
 	var err error
-	if config.Bgp.PeerIP == "" {
-		gw, err = gateway()
+	if config.PeerIP == "" {
+		gw, err := gateway()
+		if err != nil {
+			return nil, fmt.Errorf("Unable to get gw ip: %v", err)
+		}
 		c.peerIP = gw
 	} else {
-		c.peerIP = net.ParseIP(config.Bgp.PeerIP)
+		c.peerIP = net.ParseIP(config.PeerIP)
+	}
+	if config.LocalIP == "" {
 		gw, err = via(c.peerIP)
+		if err != nil {
+			return nil, fmt.Errorf("Unable to get gw ip: %v", err)
+		}
+		c.localIP, err = localAddress(gw)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		c.localIP = net.ParseIP(config.LocalIP)
 	}
-	if err != nil || c.peerIP == nil {
-		return nil, fmt.Errorf("Unable to get peer IP : %v", err)
-	}
-	c.communities = config.Bgp.Communities
-	switch config.Bgp.Origin {
+	c.communities = config.Communities
+	switch config.Origin {
 	case "igp":
 		c.origin = 0
 	case "egp":
@@ -47,24 +59,19 @@ func NewController(config *c.Config) (*Controller, error) {
 	}
 	s := gobgp.NewBgpServer()
 	go s.Serve()
-	localAddr, err := localAddress(gw)
-	if err != nil {
-		return nil, err
-	}
-	c.localIP = localAddr
 	if err := s.StartBgp(context.Background(), &api.StartBgpRequest{
 		Global: &api.Global{
-			As:         uint32(config.Bgp.LocalAS),
-			RouterId:   localAddr.String(),
+			As:         uint32(config.LocalAS),
+			RouterId:   c.localIP.String(),
 			ListenPort: -1, // gobgp won't listen on tcp:179
 		},
 	}); err != nil {
 		return nil, fmt.Errorf("Unable to start bgp: %v", err)
 	}
 	c.s = s
-	c.peerAS = config.Bgp.PeerAS
+	c.peerAS = config.PeerAS
 	// set mh by default for all ebgp peers
-	if c.peerAS != config.Bgp.LocalAS {
+	if c.peerAS != config.LocalAS {
 		c.multiHop = true
 	}
 	return c, nil
@@ -108,23 +115,21 @@ func (c *Controller) getApiPath(route *net.IPNet) *api.Path {
 	})
 	attrs := []*any.Any{a1, a2, a3}
 	return &api.Path{
-		Family:    &api.Family{Afi: afi, Safi: api.Family_SAFI_UNICAST},
-		AnyNlri:   nlri,
-		AnyPattrs: attrs,
+		Family: &api.Family{Afi: afi, Safi: api.Family_SAFI_UNICAST},
+		Nlri:   nlri,
+		Pattrs: attrs,
 	}
 }
 
 func (c *Controller) Announce(route *net.IPNet) error {
-	peers, err := c.s.ListPeer(context.Background(), &api.ListPeerRequest{})
-	if err != nil {
-		return err
-	}
 	var found bool
-	for _, p := range peers {
+	err := c.s.ListPeer(context.Background(), &api.ListPeerRequest{}, func(p *api.Peer) {
 		if p.Conf.NeighborAddress == c.peerIP.String() {
 			found = true
-			break
 		}
+	})
+	if err != nil {
+		return err
 	}
 	if !found {
 		if err := c.AddPeer(c.peerIP.String()); err != nil {
@@ -140,16 +145,16 @@ func (c *Controller) Withdraw(route *net.IPNet) error {
 }
 
 func (c *Controller) PeerInfo() (*api.Peer, error) {
-	peers, err := c.s.ListPeer(context.Background(), &api.ListPeerRequest{})
+	var peer *api.Peer
+	err := c.s.ListPeer(context.Background(), &api.ListPeerRequest{}, func(p *api.Peer) {
+		if p.Conf.NeighborAddress == c.peerIP.String() {
+			peer = p
+		}
+	})
 	if err != nil {
 		return nil, err
 	}
-	for _, p := range peers {
-		if p.Conf.NeighborAddress == c.peerIP.String() {
-			return p, nil
-		}
-	}
-	return nil, nil
+	return peer, nil
 }
 
 func (c *Controller) Shutdown() error {

controller/bgp_test.go

@@ -0,0 +1,95 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+
+	"github.com/golang/protobuf/ptypes"
+	"github.com/mayuresh82/gocast/config"
+	api "github.com/osrg/gobgp/api"
+	gobgp "github.com/osrg/gobgp/pkg/server"
+	"github.com/stretchr/testify/assert"
+)
+
+type BgpListener struct {
+	s          *gobgp.BgpServer
+	recvdPaths chan string
+}
+
+var listener *BgpListener
+
+// NewBgpListener starts a local BGP server for testing purposes
+func NewBgpListener(localAS int) (*BgpListener, error) {
+	s := gobgp.NewBgpServer()
+	go s.Serve()
+	if err := s.StartBgp(context.Background(), &api.StartBgpRequest{
+		Global: &api.Global{
+			As:       uint32(localAS),
+			RouterId: "100.100.100.100",
+		},
+	}); err != nil {
+		return nil, fmt.Errorf("Unable to start bgp: %v", err)
+	}
+	n := &BgpListener{s: s, recvdPaths: make(chan string)}
+	err := s.MonitorTable(context.Background(), &api.MonitorTableRequest{TableType: api.TableType_ADJ_IN}, func(p *api.Path) {
+		// assumes v4 only paths !
+		var value ptypes.DynamicAny
+		if err := ptypes.UnmarshalAny(p.Nlri, &value); err != nil {
+			return
+		}
+		nlri := value.Message.(*api.IPAddressPrefix)
+		n.recvdPaths <- fmt.Sprintf("%s/%d", nlri.Prefix, nlri.PrefixLen)
+	})
+	if err != nil {
+		return nil, err
+	}
+	if err := s.AddPeer(context.Background(), &api.AddPeerRequest{
+		Peer: &api.Peer{
+			Conf: &api.PeerConf{
+				NeighborAddress: "127.0.0.1",
+				PeerAs:          11111,
+			},
+			Transport: &api.Transport{PassiveMode: true},
+		},
+	}); err != nil {
+		return nil, err
+	}
+	return n, nil
+}
+
+func (l *BgpListener) Shutdown() error {
+	if err := l.s.StopBgp(context.Background(), &api.StopBgpRequest{}); err != nil {
+		return err
+	}
+	return nil
+}
+
+// This test tests the BGP controller talking to a local BGP
+// listener. It needs a few seconds to pass and *may* time out
+// if the test timeouts are very small. It also needs to be run as
+// root (sudo)
+func TestBgpNew(t *testing.T) {
+	a := assert.New(t)
+	c := config.BgpConfig{
+		LocalAS:     11111,
+		PeerAS:      22222,
+		PeerIP:      "127.0.0.1",
+		LocalIP:     "192.168.1.100",
+		Communities: []string{"100:100"},
+		Origin:      "igp",
+	}
+	ctrl, err := NewController(c)
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	_, ipnet, _ := net.ParseCIDR("20.30.40.0/24")
+	if err := ctrl.Announce(ipnet); err != nil {
+		a.FailNow(err.Error())
+	}
+
+	path := <-listener.recvdPaths
+	a.Equal("20.30.40.0/24", path)
+	ctrl.Shutdown()
+}

controller/consul.go

@@ -3,12 +3,13 @@ package controller
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/golang/glog"
 	"net/http"
 	"net/url"
 	"os"
 	"strings"
 	"time"
+
+	"github.com/golang/glog"
 )
 
 const (
@@ -20,10 +21,18 @@ const (
 	localHealthCheckurl  = "/agent/checks"
 )
 
+type Clienter interface {
+	Get(url string) (*http.Response, error)
+}
+
+type Client struct {
+	*http.Client
+}
+
 type ConsulMon struct {
 	addr   string
 	node   string
-	client *http.Client
+	client Clienter
 }
 
 type ConsulServiceData struct {
@@ -167,7 +176,7 @@ func (c *ConsulMon) healthCheckRemote(service string) (bool, error) {
 // healthCheck determines if we should use the local agent
 // If the address contains "localhost", then it presumes that the local agent is to be used.
 func (c *ConsulMon) healthCheck(service string) (bool, error) {
-	usingLocalAgent := strings.Contains(c.addr, "localhost")
+	usingLocalAgent := strings.Contains(c.addr, "localhost") || strings.Contains(c.addr, "127.0.0.1")
 	if usingLocalAgent {
 		return c.healthCheckLocal(service)
 	}

controller/consul_test.go

@@ -0,0 +1,170 @@
+package controller
+
+import (
+	"bytes"
+	"io/ioutil"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var mockConsulData = map[string]string{
+	"single-app": `{"Services": {
+		"test-app-1": {
+			"ID": "test-app-1",
+			"Service": "test-service",
+			"Tags": [
+				"enable_gocast", "gocast_vip=1.1.1.1/32", "gocast_monitor=consul"
+			]
+		}
+	}}`,
+	"single-app-no-match": `{"Services": {
+		"test-app-1": {
+			"ID": "test-app-1",
+			"Service": "test-service",
+			"Tags": [
+				"foo"
+			]
+		}
+	}}`,
+	"single-app-no-vip": `{"Services": {
+		"test-app-1": {
+			"ID": "test-app-1",
+			"Service": "test-service",
+			"Tags": [
+				"enable_gocast", "gocast_monitor=consul"
+			]
+		}
+	}}`,
+}
+
+var mockConsulCheckData = map[string]string{
+	"remote-pass": `[
+		{
+		"Node": "test-node1",
+		"Status": "passing",
+		"ServiceName": "test-service"
+		},
+		{
+		"Node": "test-node2",
+		"Status": "passing",
+		"ServiceName": "test-service"
+		}
+	]`,
+	"remote-fail": `[
+		{
+		"Node": "test-node1",
+		"Status": "failed",
+		"ServiceName": "test-service"
+		}
+	]`,
+	"local-pass": `{
+  		"service:test-service": {
+  		  "Node": "test-node1",
+  		  "Status": "passing",
+  		  "ServiceName": "test-service"
+  		}
+	}`,
+	"local-fail": `{
+  		"service:test-service": {
+  		  "Node": "test-node1",
+  		  "Status": "failed",
+  		  "ServiceName": "test-service"
+  		}
+	}`,
+}
+
+type MockClient struct {
+	get func(url string) (*http.Response, error)
+}
+
+func (c *MockClient) Get(url string) (*http.Response, error) {
+	if c.get != nil {
+		return c.get(url)
+	}
+	return nil, nil
+}
+
+func TestQueryServices(t *testing.T) {
+	a := assert.New(t)
+	client := &MockClient{}
+	cm := &ConsulMon{
+		addr: "foo", node: "test", client: client,
+	}
+
+	// test valid app
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulData["single-app"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	apps, err := cm.queryServices()
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.Equal(1, len(apps))
+	app, _ := NewApp("test-service", "1.1.1.1/32", []string{"consul"}, []string{}, "consul")
+	a.True(app.Equal(apps[0]))
+
+	// test no match
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulData["single-app-no-match"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	apps, err = cm.queryServices()
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.Equal(0, len(apps))
+
+	// test missing vip
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulData["single-app-no-vip"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	apps, _ = cm.queryServices()
+	a.Equal(0, len(apps))
+}
+
+func TestHealthCheck(t *testing.T) {
+	a := assert.New(t)
+	client := &MockClient{}
+	cm := &ConsulMon{node: "test-node1", client: client}
+
+	// test remote checks
+	cm.addr = "http://remote/check"
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulCheckData["remote-pass"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	check, err := cm.healthCheck("test-service")
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.True(check)
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulCheckData["remote-fail"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	check, _ = cm.healthCheck("test-service")
+	a.False(check)
+
+	// test local checks
+	cm.addr = "http://localhost/check"
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulCheckData["local-pass"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	check, _ = cm.healthCheck("test-service")
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.True(check)
+	cm.addr = "http://127.0.0.1/check"
+	client.get = func(url string) (*http.Response, error) {
+		b := bytes.NewBuffer([]byte(mockConsulCheckData["local-fail"]))
+		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
+	}
+	check, _ = cm.healthCheck("test-service")
+	a.False(check)
+}

controller/monitor.go

@@ -2,14 +2,15 @@ package controller
 
 import (
 	"fmt"
-	"github.com/golang/glog"
-	c "github.com/mayuresh82/gocast/config"
-	api "github.com/osrg/gobgp/api"
 	"net"
 	"os/exec"
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/golang/glog"
+	c "github.com/mayuresh82/gocast/config"
+	api "github.com/osrg/gobgp/api"
 )
 
 const (
@@ -52,6 +53,7 @@ func execMonitor(cmd string) bool {
 	return true
 }
 
+// appMon maintains the state of a registered app
 type appMon struct {
 	app       *App
 	done      chan bool
@@ -59,6 +61,7 @@ type appMon struct {
 	checkOn   bool
 }
 
+// MonitorMgr manages the lifecycle of registered apps
 type MonitorMgr struct {
 	monitors map[string]*appMon
 	cleanups map[string]chan bool
@@ -70,7 +73,7 @@ type MonitorMgr struct {
 }
 
 func NewMonitor(config *c.Config) *MonitorMgr {
-	ctrl, err := NewController(config)
+	ctrl, err := NewController(config.Bgp)
 	if err != nil {
 		glog.Exitf("Failed to start BGP controller: %v", err)
 	}
@@ -107,6 +110,8 @@ func NewMonitor(config *c.Config) *MonitorMgr {
 	return mon
 }
 
+// consulMon periodically queries consul for apps that need to be
+// registered and adds them to the monitor manager
 func (m *MonitorMgr) consulMon() {
 	for {
 		apps, err := m.consul.queryServices()
@@ -144,6 +149,7 @@ func (m *MonitorMgr) consulMon() {
 	}
 }
 
+// Add adds a new app into monitor manager
 func (m *MonitorMgr) Add(app *App) {
 	// check if already running
 	m.Lock()
@@ -165,6 +171,8 @@ func (m *MonitorMgr) Add(app *App) {
 	glog.Infof("Registered a new app: %v", app)
 }
 
+// Remove removes an app from monitor manager, stops BGP
+/// announcement and cleans up state
 func (m *MonitorMgr) Remove(appName string) {
 	if a, ok := m.monitors[appName]; ok {
 		if a.checkOn {
@@ -254,6 +262,8 @@ func (m *MonitorMgr) checkCond(am *appMon) error {
 	return nil
 }
 
+// runLoop periodically checks if an app passes healthchecks
+// and needs VIP announcement
 func (m *MonitorMgr) runLoop(am *appMon) {
 	am.checkOn = true
 	if err := m.checkCond(am); err != nil {
@@ -274,6 +284,7 @@ func (m *MonitorMgr) runLoop(am *appMon) {
 	}
 }
 
+// CloseAll shuts down all BGP sessions removes state
 func (m *MonitorMgr) CloseAll() {
 	glog.Infof("Shutting down all open bgp sessions")
 	if err := m.ctrl.Shutdown(); err != nil {
@@ -294,6 +305,7 @@ func (m *MonitorMgr) CloseAll() {
 	}
 }
 
+// CleanUp periodically monitors for stale apps and cleans them up
 func (m *MonitorMgr) Cleanup(app string, exit chan bool) {
 	t := time.NewTimer(m.config.Agent.CleanupTimer)
 	defer t.Stop()
@@ -310,6 +322,7 @@ func (m *MonitorMgr) Cleanup(app string, exit chan bool) {
 	}
 }
 
+// GetInfo returns basic BGP info for established peers
 func (m *MonitorMgr) GetInfo() (*api.Peer, error) {
 	return m.ctrl.PeerInfo()
 }

controller/monitor_test.go

@@ -0,0 +1,35 @@
+package controller
+
+import (
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPortMonitor(t *testing.T) {
+	a := assert.New(t)
+	addr, _ := net.ResolveTCPAddr("tcp", ":33333")
+	conn, err := net.ListenTCP("tcp", addr)
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.True(portMonitor("tcp", "33333"))
+	a.False(portMonitor("tcp", "44444"))
+	conn.Close()
+
+	uaddr, _ := net.ResolveUDPAddr("udp", ":33333")
+	udpconn, err := net.ListenUDP("udp", uaddr)
+	if err != nil {
+		a.FailNow(err.Error())
+	}
+	a.True(portMonitor("udp", "33333"))
+	a.False(portMonitor("udp", "44444"))
+	udpconn.Close()
+}
+
+func TestExecMonitor(t *testing.T) {
+	a := assert.New(t)
+	a.True(execMonitor("echo foo"))
+	a.False(execMonitor("echo foo && false"))
+}

controller/system.go

@@ -7,9 +7,21 @@ import (
 	"strings"
 )
 
+var execCmd = "bash"
+
+func getCmdList(mainCmd string) []string {
+	cmdList := []string{}
+	if execCmd == "bash" {
+		cmdList = append(cmdList, "-c")
+	}
+	cmdList = append(cmdList, mainCmd)
+	return cmdList
+}
+
 func gateway() (net.IP, error) {
 	cmd := `ip route | grep "^default" | cut -d" " -f3`
-	out, err := exec.Command("bash", "-c", cmd).Output()
+	cmdList := getCmdList(cmd)
+	out, err := exec.Command(execCmd, cmdList...).Output()
 	if err != nil {
 		return nil, fmt.Errorf("Failed to execute command: %s", cmd)
 	}
@@ -18,7 +30,8 @@ func gateway() (net.IP, error) {
 
 func via(dest net.IP) (net.IP, error) {
 	cmd := fmt.Sprintf(`ip route get %s | grep via | cut -d" " -f3`, dest.String())
-	out, err := exec.Command("bash", "-c", cmd).Output()
+	cmdList := getCmdList(cmd)
+	out, err := exec.Command(execCmd, cmdList...).Output()
 	if err != nil {
 		return nil, fmt.Errorf("Failed to execute command: %s", cmd)
 	}
@@ -54,7 +67,8 @@ func addLoopback(name string, addr *net.IPNet) error {
 		label = label[:15]
 	}
 	cmd := fmt.Sprintf("ip address add %s/%d dev lo label %s", addr.IP.String(), prefixLen, label)
-	_, err := exec.Command("bash", "-c", cmd).Output()
+	cmdList := getCmdList(cmd)
+	_, err := exec.Command(execCmd, cmdList...).Output()
 	if err != nil {
 		return fmt.Errorf("Failed to Add loopback command: %s: %v", cmd, err)
 	}
@@ -64,7 +78,8 @@ func addLoopback(name string, addr *net.IPNet) error {
 func deleteLoopback(addr *net.IPNet) error {
 	prefixLen, _ := addr.Mask.Size()
 	cmd := fmt.Sprintf("ip address delete %s/%d dev lo", addr.IP.String(), prefixLen)
-	_, err := exec.Command("bash", "-c", cmd).Output()
+	cmdList := getCmdList(cmd)
+	_, err := exec.Command(execCmd, cmdList...).Output()
 	if err != nil {
 		return fmt.Errorf("Failed to delete loopback command: %s: %v", cmd, err)
 	}
@@ -76,7 +91,8 @@ func natRule(op string, vip, localAddr net.IP, protocol, port string) error {
 		"iptables -t nat -%s PREROUTING -p %s -d %s --dport %s -j DNAT --to-destination %s:%s",
 		op, protocol, vip.String(), port, localAddr.String(), port,
 	)
-	_, err := exec.Command("bash", "-c", cmd).Output()
+	cmdList := getCmdList(cmd)
+	_, err := exec.Command(execCmd, cmdList...).Output()
 	if err != nil {
 		return fmt.Errorf("Failed to %s nat rule: %s: %v", op, cmd, err)
 	}

controller/system_test.go

@@ -0,0 +1,70 @@
+package controller
+
+import (
+	"fmt"
+	"net"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGateway(t *testing.T) {
+	execCmd = os.Args[0]
+	os.Setenv("test_name", "test_gateway")
+	gw, err := gateway()
+	assert.Nil(t, err)
+	assert.Equal(t, "10.1.1.1", gw.String())
+}
+
+func TestVia(t *testing.T) {
+	execCmd = os.Args[0]
+	os.Setenv("test_name", "test_via")
+	ip, err := via(net.ParseIP("10.1.2.100"))
+	assert.Nil(t, err)
+	assert.Equal(t, "10.1.2.1", ip.String())
+
+	os.Setenv("test_name", "test_via_none")
+	ip, err = via(net.ParseIP("10.1.4.1"))
+	assert.Nil(t, err)
+	assert.Equal(t, "10.1.4.1", ip.String())
+}
+
+func TestAddLoopback(t *testing.T) {
+	execCmd = os.Args[0]
+	os.Setenv("test_name", "test_add_pass")
+	_, ipnet, _ := net.ParseCIDR("1.1.1.1/32")
+	err := addLoopback("test_app", ipnet)
+	assert.Nil(t, err)
+
+	os.Setenv("test_name", "test_add_fail")
+	_, ipnet, _ = net.ParseCIDR("1.1.1.1/32")
+	err = addLoopback("test_app", ipnet)
+	assert.NotNil(t, err)
+}
+
+func TestMain(m *testing.M) {
+	switch os.Getenv("test_name") {
+	case "test_gateway":
+		fmt.Println("10.1.1.1")
+	case "test_via":
+		fmt.Println("10.1.2.1")
+	case "test_via_none":
+		break
+	case "test_add_fail":
+		os.Exit(1)
+	default:
+		fmt.Println("success")
+	}
+	if os.Getenv("test_name") != "" {
+		return
+	}
+	var err error
+	listener, err = NewBgpListener(22222)
+	if err != nil {
+		panic(err)
+	}
+	code := m.Run()
+	listener.Shutdown()
+	os.Exit(code)
+}