Test: Add readyz tests

cmd/readyz.go

@@ -1,23 +1,19 @@
 package cmd
 
 import (
-	"crypto/tls"
 	"fmt"
-	"net/http"
 	"os"
-	"path"
-	"strings"
 	"time"
 
 	"github.com/axllent/mailpit/config"
+	"github.com/axllent/mailpit/internal/healthcheck"
 	"github.com/spf13/cobra"
 )
 
 var (
-	useHTTPS        bool
-	readyzWait      bool
-	readyzTimeout   time.Duration
-	readyzPollEvery = time.Second
+	useHTTPS      bool
+	readyzWait    bool
+	readyzTimeout time.Duration
 )
 
 // readyzCmd represents the healthcheck command
@@ -25,76 +21,30 @@ var readyzCmd = &cobra.Command{
 	Use:   "readyz",
 	Short: "Run a healthcheck to test if Mailpit is running",
 	Long: `This command connects to the /readyz endpoint of a running Mailpit server
-and exits with a status of 0 if the connection is successful, else with a 
+and exits with a status of 0 if the connection is successful, else with a
 status 1 if unhealthy.
 
 If running within Docker, it should automatically detect environment
 settings to determine the HTTP bind interface & port.
 `,
 	Run: func(_ *cobra.Command, _ []string) {
-		webroot := strings.TrimRight(path.Join("/", config.Webroot, "/"), "/") + "/"
-		proto := "http"
-		if useHTTPS {
-			proto = "https"
-		}
-
-		uri := fmt.Sprintf("%s://%s%sreadyz", proto, config.HTTPListen, webroot)
-
-		conf := &http.Transport{
-			IdleConnTimeout:       time.Second * 5,
-			ExpectContinueTimeout: time.Second * 5,
-			TLSHandshakeTimeout:   time.Second * 5,
-			// do not verify TLS if this instance is using HTTPS as we connect using IP
-			// so won't be the same as the cert
-			TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // #nosec
-		}
-		client := &http.Client{Transport: conf}
+		uri := healthcheck.URI(config.HTTPListen, config.Webroot, useHTTPS)
+		client := healthcheck.NewClient()
 
+		var err error
 		if readyzWait {
-			if err := waitForReady(client, uri, readyzTimeout); err != nil {
-				fmt.Fprintln(os.Stderr, err)
-				os.Exit(1)
-			}
-			return
+			err = healthcheck.Wait(client, uri, readyzTimeout)
+		} else {
+			err = healthcheck.Check(client, uri)
 		}
 
-		if err := checkReady(client, uri); err != nil {
+		if err != nil {
 			fmt.Fprintln(os.Stderr, err)
 			os.Exit(1)
 		}
 	},
 }
 
-func checkReady(client *http.Client, uri string) error {
-	res, err := client.Get(uri)
-	if err != nil {
-		return err
-	}
-	defer res.Body.Close()
-
-	if res.StatusCode != http.StatusOK {
-		return fmt.Errorf("unexpected status: %s", res.Status)
-	}
-
-	return nil
-}
-
-func waitForReady(client *http.Client, uri string, timeout time.Duration) error {
-	deadline := time.Now().Add(timeout)
-
-	for {
-		if err := checkReady(client, uri); err == nil {
-			return nil
-		}
-
-		if time.Now().After(deadline) {
-			return fmt.Errorf("timed out after %s waiting for Mailpit to become ready", timeout)
-		}
-
-		time.Sleep(readyzPollEvery)
-	}
-}
-
 func init() {
 	rootCmd.AddCommand(readyzCmd)
 

cmd/readyz_test.go

@@ -1,65 +0,0 @@
-package cmd
-
-import (
-	"net/http"
-	"net/http/httptest"
-	"testing"
-	"time"
-)
-
-func TestCheckReady(t *testing.T) {
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-	}))
-	t.Cleanup(srv.Close)
-
-	if err := checkReady(srv.Client(), srv.URL); err != nil {
-		t.Fatalf("checkReady() error = %v", err)
-	}
-}
-
-func TestWaitForReadyRetriesUntilSuccess(t *testing.T) {
-	oldPoll := readyzPollEvery
-	readyzPollEvery = time.Millisecond
-	t.Cleanup(func() { readyzPollEvery = oldPoll })
-
-	var calls int
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		calls++
-		if calls == 1 {
-			w.WriteHeader(http.StatusServiceUnavailable)
-			return
-		}
-		w.WriteHeader(http.StatusOK)
-	}))
-	t.Cleanup(srv.Close)
-
-	if err := waitForReady(srv.Client(), srv.URL, 100*time.Millisecond); err != nil {
-		t.Fatalf("waitForReady() error = %v", err)
-	}
-
-	if calls < 2 {
-		t.Fatalf("waitForReady() calls = %d, want at least 2", calls)
-	}
-}
-
-func TestWaitForReadyTimesOut(t *testing.T) {
-	oldPoll := readyzPollEvery
-	readyzPollEvery = time.Millisecond
-	t.Cleanup(func() { readyzPollEvery = oldPoll })
-
-	var calls int
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		calls++
-		w.WriteHeader(http.StatusServiceUnavailable)
-	}))
-	t.Cleanup(srv.Close)
-
-	if err := waitForReady(srv.Client(), srv.URL, 5*time.Millisecond); err == nil {
-		t.Fatal("waitForReady() error = nil, want timeout")
-	}
-
-	if calls == 0 {
-		t.Fatal("waitForReady() did not call the endpoint")
-	}
-}