add tests

Author: Pavani-Panakanti

main.go

@@ -52,6 +52,7 @@ import (
 var (
 	scheme              = runtime.NewScheme()
 	LOCAL_IPAMD_ADDRESS = "127.0.0.1:50051"
+	npaSocketPath       = "/var/run/aws-node/npa.sock"
 )
 
 func init() {
@@ -134,9 +135,15 @@ func main() {
 
 	// CNI makes rpc calls to NP agent regardless NP is enabled or not
 	// need to start rpc always
+	// todo: add a liveness probe to this gRPC server and remove closing based on this errCh, liveness probe will check and re-start this container
+	errCh, err := rpc.RunRPCHandler(policyEndpointController, npaSocketPath)
+	if err != nil {
+		log.Errorf("Failed to set up gRPC Handler %v", err)
+		os.Exit(1)
+	}
 	go func() {
-		if err := rpc.RunRPCHandler(policyEndpointController); err != nil {
-			log.Errorf("Failed to set up gRPC Handler %v", err)
+		if err := <-errCh; err != nil {
+			log.Errorf("gRPC server stopped: %v", err)
 			os.Exit(1)
 		}
 	}()

pkg/rpc/rpc_handler.go

@@ -42,7 +42,6 @@ var (
 )
 
 const (
-	npaSocketPath         = "/var/run/aws-node/npa.sock"
 	grpcHealthServiceName = "grpc.health.v1.np-agent"
 )
 
@@ -160,18 +159,21 @@ func (s *server) DeletePodNp(ctx context.Context, in *rpc.DeleteNpRequest) (*rpc
 }
 
 // RunRPCHandler handles request from gRPC
-func RunRPCHandler(policyReconciler *controllers.PolicyEndpointsReconciler) error {
+func RunRPCHandler(policyReconciler *controllers.PolicyEndpointsReconciler, npaSocketPath string) (<-chan error, error) {
 	log().Infof("Serving RPC Handler on Unix socket: %s", npaSocketPath)
 
 	if _, err := os.Stat(npaSocketPath); err == nil {
 		log().Infof("Removing stale socket file: %s", npaSocketPath)
-		os.Remove(npaSocketPath)
+		err = os.Remove(npaSocketPath)
+		if err != nil {
+			log().Warnf("got error in removing socket file %v", err)
+		}
 	}
 
 	listener, err := net.Listen("unix", npaSocketPath)
 	if err != nil {
 		log().Errorf("Failed to listen on unix socket: %v", err)
-		return errors.Wrap(err, "network policy agent: failed to listen on unix socket")
+		return nil, errors.Wrap(err, "network policy agent: failed to listen on unix socket")
 	}
 	grpcServer := grpc.NewServer()
 	rpc.RegisterNPBackendServer(grpcServer, &server{policyReconciler: policyReconciler})
@@ -182,10 +184,13 @@ func RunRPCHandler(policyReconciler *controllers.PolicyEndpointsReconciler) erro
 
 	// Register reflection service on gRPC server.
 	reflection.Register(grpcServer)
-	if err := grpcServer.Serve(listener); err != nil {
-		log().Errorf("Failed to start server on gRPC port: %v", err)
-		return errors.Wrap(err, "network policy agent: failed to start server on gPRC port")
-	}
+	errCh := make(chan error, 1)
+	go func() {
+		if err := grpcServer.Serve(listener); err != nil {
+			errCh <- errors.Wrap(err, "network policy agent: grpc serve failed")
+		}
+		close(errCh)
+	}()
 	log().Info("Done with RPC Handler initialization")
-	return nil
+	return errCh, nil
 }

pkg/rpc/rpc_handler_test.go

@@ -0,0 +1,46 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//     http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package rpc
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRunRPCHandler_NoExistingSocket(t *testing.T) {
+	testSocketPath := "/tmp/test-rpc-handler.sock"
+	defer os.Remove(testSocketPath)
+
+	errCh, err := RunRPCHandler(nil, testSocketPath)
+	assert.Nil(t, err)
+	assert.NotNil(t, errCh)
+}
+
+func TestRunRPCHandler_StaleSocketCleanup(t *testing.T) {
+	testSocketPath := "/tmp/temp-rpc-handler.sock"
+
+	// Create a stale socket file
+	file, err := os.Create(testSocketPath)
+	if err != nil {
+		t.Fatalf("Failed to create stale socket file: %v", err)
+	}
+	file.Close()
+	defer os.Remove(testSocketPath)
+
+	errCh, err := RunRPCHandler(nil, testSocketPath)
+	assert.Nil(t, err)
+	assert.NotNil(t, errCh)
+}