From 3e7cea7bd39253718bc3d2f8b10ac5fc9b98cbc2 Mon Sep 17 00:00:00 2001
From: Alex Jones <alexsimonjones@gmail.com>
Date: Fri, 29 Dec 2023 14:15:32 +0000
Subject: [PATCH] chore: added basic server startup test (#817)

Signed-off-by: Alex Jones <alexsimonjones@gmail.com>
---
 pkg/server/server.go      |  6 ++++++
 pkg/server/server_test.go | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 pkg/server/server_test.go

diff --git a/pkg/server/server.go b/pkg/server/server.go
index 842e79c2cb..00862c5fa0 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -42,6 +42,7 @@ type Config struct {
 	Handler        *handler
 	Logger         *zap.Logger
 	metricsServer  *http.Server
+	listener       net.Listener
 }
 
 type Health struct {
@@ -56,6 +57,10 @@ var health = Health{
 	Failure: 0,
 }
 
+func (s *Config) Shutdown() error {
+	return s.listener.Close()
+}
+
 func (s *Config) Serve() error {
 
 	var lis net.Listener
@@ -65,6 +70,7 @@ func (s *Config) Serve() error {
 	if err != nil {
 		return err
 	}
+	s.listener = lis
 	s.Logger.Info(fmt.Sprintf("binding api to %s", s.Port))
 	grpcServerUnaryInterceptor := grpc.UnaryInterceptor(logInterceptor(s.Logger))
 	grpcServer := grpc.NewServer(grpcServerUnaryInterceptor)
diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go
new file mode 100644
index 0000000000..98d9743cb1
--- /dev/null
+++ b/pkg/server/server_test.go
@@ -0,0 +1,41 @@
+package server
+
+import (
+	"github.com/fatih/color"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+	"os"
+	"sync"
+	"testing"
+)
+
+func TestServerInit(t *testing.T) {
+	logger, err := zap.NewDevelopment()
+	if err != nil {
+		color.Red("failed to create logger: %v", err)
+		os.Exit(1)
+	}
+	defer logger.Sync()
+	server_config := Config{
+		Backend:     "openai",
+		Port:        "0",
+		MetricsPort: "0",
+		Token:       "none",
+		Logger:      logger,
+	}
+	var wg sync.WaitGroup
+
+	go func() {
+		wg.Add(1)
+		err := server_config.Serve()
+		if err != nil {
+			assert.Fail(t, "serve: %s", err.Error())
+		}
+		server_config.Shutdown()
+		if err != nil {
+			assert.Fail(t, "shutdown: %s", err.Error())
+		}
+		wg.Done()
+	}()
+	wg.Wait()
+}