1 files changed, 60 insertions, 11 deletions
diff --git a/internal/sshd/sshd_test.go b/internal/sshd/sshd_test.go
index d1891ec..9187140 100644
--- a/internal/sshd/sshd_test.go
+++ b/internal/sshd/sshd_test.go
@@ -4,6 +4,8 @@ import (
 	"testing"
 	"context"
 	"path"
+	"net/http/httptest"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -17,18 +19,55 @@ const serverUrl = "127.0.0.1:50000"
 func TestShutdown(t *testing.T) {
 	s := setupServer(t)
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	go func() { require.NoError(t, s.ListenAndServe(context.Background())) }()
 
-	done := make(chan bool, 1)
-	go func() {
-		require.NoError(t, s.serve(ctx))
-		done <- true
-	}()
+	verifyStatus(t, s, StatusReady)
+
+	s.wg.Add(1)
 
 	require.NoError(t, s.Shutdown())
+	verifyStatus(t, s, StatusOnShutdown)
+
+	s.wg.Done()
+
+	verifyStatus(t, s, StatusClosed)
+}
+
+func TestReadinessProbe(t *testing.T) {
+	s := &Server{Config: &config.Config{Server: config.DefaultServerConfig}}
+
+	require.Equal(t, StatusStarting, s.getStatus())
+
+	mux := s.MonitoringServeMux()
+
+	req := httptest.NewRequest("GET", "/start", nil)
+
+	r := httptest.NewRecorder()
+	mux.ServeHTTP(r, req)
+	require.Equal(t, 503, r.Result().StatusCode)
+
+	s.changeStatus(StatusReady)
 
-	require.True(t, <-done, "the accepting loop must be interrupted")
+	r = httptest.NewRecorder()
+	mux.ServeHTTP(r, req)
+	require.Equal(t, 200, r.Result().StatusCode)
+
+	s.changeStatus(StatusOnShutdown)
+
+	r = httptest.NewRecorder()
+	mux.ServeHTTP(r, req)
+	require.Equal(t, 503, r.Result().StatusCode)
+}
+
+func TestLivenessProbe(t *testing.T) {
+	s := &Server{Config: &config.Config{Server: config.DefaultServerConfig}}
+	mux := s.MonitoringServeMux()
+
+	req := httptest.NewRequest("GET", "/health", nil)
+
+	r := httptest.NewRecorder()
+	mux.ServeHTTP(r, req)
+	require.Equal(t, 200, r.Result().StatusCode)
 }
 
 func setupServer(t *testing.T) *Server {
@@ -42,8 +81,18 @@ func setupServer(t *testing.T) *Server {
 
 	cfg := &config.Config{RootDir: "/tmp", GitlabUrl: url, Server: srvCfg}
 
-	s := &Server{Config: cfg}
-	require.NoError(t, s.listen())
+	return &Server{Config: cfg}
+}
+
+func verifyStatus(t *testing.T, s *Server, st status) {
+	for i := 5; i < 500; i+=50 {
+		if s.getStatus() == st {
+			break
+		}
+
+		// Sleep incrementally ~2s in total
+		time.Sleep(time.Duration(i) * time.Millisecond)
+	}
 
-	return s
+	require.Equal(t, s.getStatus(), st)
 }