fixing tests

Author: Chao Xu

pkg/agent/agentclient/client_test.go

@@ -17,14 +17,15 @@ import (
 func TestServeData_HTTP(t *testing.T) {
 	var err error
 	var stream agent.AgentService_ConnectClient
+	stopCh := make(chan struct{})
 	client := &AgentClient{
 		connContext: make(map[int64]*connContext),
+		stopCh:      stopCh,
 	}
 	client.stream, stream = pipe2()
-	stopCh := make(chan struct{})
 
 	// Start agent
-	go client.Serve(stopCh)
+	go client.Serve()
 	defer close(stopCh)
 
 	// Start test http server as remote service
@@ -112,14 +113,15 @@ func TestServeData_HTTP(t *testing.T) {
 
 func TestClose_Client(t *testing.T) {
 	var stream agent.AgentService_ConnectClient
+	stopCh := make(chan struct{})
 	client := &AgentClient{
 		connContext: make(map[int64]*connContext),
+		stopCh:      stopCh,
 	}
 	client.stream, stream = pipe2()
-	stopCh := make(chan struct{})
 
 	// Start agent
-	go client.Serve(stopCh)
+	go client.Serve()
 	defer close(stopCh)
 
 	// Start test http server as remote service

pkg/agent/agentclient/stream.go

@@ -109,8 +109,11 @@ func (c *RedialableAgentClient) probe() {
 		case <-c.stopCh:
 			return
 		case <-time.After(c.Interval):
+			if c.conn == nil {
+				continue
+			}
 			// health check
-			if c.conn != nil && c.conn.GetState() == connectivity.Ready {
+			if c.conn.GetState() == connectivity.Ready {
 				continue
 			} else {
 				klog.Infof("Connection state %v", c.conn.GetState())
@@ -256,7 +259,7 @@ func (c *RedialableAgentClient) reconnect() {
 		}
 		switch {
 		case r.serverID == c.serverID:
-			klog.Info("reconnected to %s", serverID)
+			klog.Infof("reconnected to %s", r.serverID)
 			c.conn = r.grpcConn
 			c.stream = r.agentServiceClient
 			c.doneReconnect(nil)
@@ -308,6 +311,7 @@ func serverCount(stream agent.AgentService_ConnectClient) (int, error) {
 }
 
 func serverID(stream agent.AgentService_ConnectClient) (string, error) {
+	// TODO: this is a blocking call. Add a timeout?
 	md, err := stream.Header()
 	if err != nil {
 		return "", err

pkg/agent/agentclient/stream_test.go

@@ -7,9 +7,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 	"k8s.io/klog"
 	"sigs.k8s.io/apiserver-network-proxy/proto/agent"
+	"sigs.k8s.io/apiserver-network-proxy/proto/header"
 )
 
 func TestReconnectExits(t *testing.T) {
@@ -19,7 +22,7 @@ func TestReconnectExits(t *testing.T) {
 
 	time.Sleep(time.Millisecond)
 
-	client, err := NewRedialableAgentClient("localhost:8899", grpc.WithInsecure())
+	client, err := NewRedialableAgentClient("localhost:8899", uuid.New().String(), &ClientSet{}, grpc.WithInsecure())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -83,6 +86,11 @@ func newTestServer(addr string) *testServer {
 func (s *testServer) Connect(stream agent.AgentService_ConnectServer) error {
 	stopCh := make(chan error)
 
+	h := metadata.Pairs(header.ServerID, "", header.ServerCount, "1")
+	if err := stream.SendHeader(h); err != nil {
+		return err
+	}
+
 	// Recv only
 	go func() {
 		for {

tests/proxy_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"google.golang.org/grpc"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/agent/agentclient"
 	"sigs.k8s.io/apiserver-network-proxy/pkg/agent/agentserver"
@@ -227,7 +228,7 @@ func runGRPCProxyServer() (proxy, func(), error) {
 	var err error
 	var lis, lis2 net.Listener
 
-	server := agentserver.NewProxyServer()
+	server := agentserver.NewProxyServer(uuid.New().String(), 0)
 	grpcServer := grpc.NewServer()
 	agentServer := grpc.NewServer()
 	cleanup := func() {
@@ -262,7 +263,7 @@ func runGRPCProxyServer() (proxy, func(), error) {
 
 func runHTTPConnProxyServer() (proxy, func(), error) {
 	var proxy proxy
-	server := agentserver.NewProxyServer()
+	server := agentserver.NewProxyServer(uuid.New().String(), 0)
 	agentServer := grpc.NewServer()
 
 	agent.RegisterAgentServiceServer(agentServer, server)
@@ -304,12 +305,8 @@ func runHTTPConnProxyServer() (proxy, func(), error) {
 }
 
 func runAgent(addr string, stopCh <-chan struct{}) error {
-	client, err := agentclient.NewAgentClient(addr, grpc.WithInsecure())
-	if err != nil {
-		return err
-	}
-
-	go client.Serve(stopCh)
+	client := agentclient.NewAgentClientSet(addr, uuid.New().String(), 5*time.Second, grpc.WithInsecure())
+	go client.Serve()
 
 	return nil
 }