From d104c904456b8c9d4edd53a1b5ccb20501569c77 Mon Sep 17 00:00:00 2001
From: Marc Di Luzio <marc.diluzio@gmail.com>
Date: Sun, 31 May 2020 19:48:43 +0100
Subject: [PATCH] Fix tests with a proper start-up and shutdown procedure with
 a sync

---
 pkg/server/router.go      |  3 ---
 pkg/server/router_test.go |  4 ----
 pkg/server/server.go      | 26 +++++++++++++++++++++++++-
 script/test.sh            | 15 +++++----------
 4 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/pkg/server/router.go b/pkg/server/router.go
index bb8fc78..842d542 100644
--- a/pkg/server/router.go
+++ b/pkg/server/router.go
@@ -5,14 +5,11 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/gorilla/mux"
 	"github.com/mdiluz/rove/pkg/accounts"
 )
 
 // NewRouter sets up the server mux
 func (s *Server) SetUpRouter() {
-	s.router = mux.NewRouter().StrictSlash(true)
-
 	// Set up the handlers
 	s.router.HandleFunc("/status", s.HandleStatus)
 	s.router.HandleFunc("/register", s.HandleRegister)
diff --git a/pkg/server/router_test.go b/pkg/server/router_test.go
index 6daecbd..d92c4b1 100644
--- a/pkg/server/router_test.go
+++ b/pkg/server/router_test.go
@@ -13,8 +13,6 @@ func TestHandleStatus(t *testing.T) {
 	response := httptest.NewRecorder()
 
 	s := NewServer()
-	s.Initialise()
-
 	s.HandleStatus(response, request)
 
 	var status StatusResponse
@@ -36,8 +34,6 @@ func TestHandleRegister(t *testing.T) {
 	response := httptest.NewRecorder()
 
 	s := NewServer()
-	s.Initialise()
-
 	s.HandleRegister(response, request)
 
 	var status RegisterResponse
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 8e581e5..c667855 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -1,9 +1,12 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
+	"sync"
+	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/mdiluz/rove/pkg/accounts"
@@ -25,9 +28,12 @@ type Server struct {
 	accountant *accounts.Accountant
 	world      *game.World
 
+	server *http.Server
 	router *mux.Router
 
 	persistence int
+
+	sync sync.WaitGroup
 }
 
 // ServerOption defines a server creation option
@@ -50,12 +56,16 @@ func OptionPersistentData() ServerOption {
 // NewServer sets up a new server
 func NewServer(opts ...ServerOption) *Server {
 
+	router := mux.NewRouter().StrictSlash(true)
+
 	// Set up the default server
 	s := &Server{
 		port:        8080,
 		accountant:  accounts.NewAccountant(),
 		world:       game.NewWorld(),
 		persistence: EphemeralData,
+		router:      router,
+		server:      &http.Server{Addr: ":8080", Handler: router},
 	}
 
 	// Apply all options
@@ -83,20 +93,34 @@ func (s *Server) Initialise() error {
 	s.SetUpRouter()
 	fmt.Printf("Routes Created\n")
 
+	// Add to our sync
+	s.sync.Add(1)
+
 	return nil
 }
 
 // Run executes the server
 func (s *Server) Run() {
+	defer s.sync.Done()
+
 	// Listen and serve the http requests
 	fmt.Println("Serving HTTP")
-	if err := http.ListenAndServe(fmt.Sprintf(":%d", s.port), s.router); err != nil {
+	if err := s.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 		log.Fatal(err)
 	}
 }
 
 // Close closes up the server
 func (s *Server) Close() error {
+	// Try and shut down the http server
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := s.server.Shutdown(ctx); err != nil {
+		return err
+	}
+
+	// Wait until the server is shut down
+	s.sync.Wait()
 
 	// Save the accounts if requested
 	if s.persistence == PersistentData {
diff --git a/script/test.sh b/script/test.sh
index 6fb0aee..8bd2620 100755
--- a/script/test.sh
+++ b/script/test.sh
@@ -4,18 +4,13 @@ cd "$(dirname "$0")"
 cd ..
 set -x
 
-# Test the build
-go build -v .
-
-# Run unit tests
-go test -v ./... -cover
-
-# Verify docker build
-docker build .
-
-# Run the integration tests with docker-compose
+# Build and start the service
 docker-compose up --build --detach
+
+# Run tests, including integration tests
 go test -v ./... -tags integration -cover -coverprofile=/tmp/c.out
+
+# Take down the service
 docker-compose down
 
 # Convert the coverage data to html