From 27ba2ab7648853e110c37955dcae7771ff58684a Mon Sep 17 00:00:00 2001
From: Johnny Steenbergen <johnny.steenbergen@crowdstrike.com>
Date: Tue, 18 Jun 2024 17:52:54 -0500
Subject: [PATCH] chore: add TestRun test helper for testing Run integrations

---
 runner_http.go        | 72 +++++++++++++++++++++++--------------------
 runner_http_direct.go | 64 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 33 deletions(-)
 create mode 100644 runner_http_direct.go

diff --git a/runner_http.go b/runner_http.go
index a17aef6..762de4d 100644
--- a/runner_http.go
+++ b/runner_http.go
@@ -23,6 +23,31 @@ const (
 type runnerHTTP struct{}
 
 func (r *runnerHTTP) Run(ctx context.Context, logger *slog.Logger, h Handler) {
+	s := &http.Server{
+		Addr:           fmt.Sprintf(":%d", port()),
+		Handler:        newHTTPRunMux(ctx, logger, h),
+		MaxHeaderBytes: mb,
+	}
+	go func() {
+		select {
+		case <-ctx.Done():
+			shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel()
+
+			logger.Info("shutting down HTTP server...")
+			if err := s.Shutdown(shutdownCtx); err != nil {
+				logger.Error("failed to shutdown server", "err", err)
+			}
+		}
+	}()
+
+	logger.Info("serving HTTP server on port " + strconv.Itoa(port()))
+	if err := s.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		logger.Error("unexpected shutdown of server", "err", err)
+	}
+}
+
+func newHTTPRunMux(ctx context.Context, logger *slog.Logger, h Handler) *http.ServeMux {
 	mux := http.NewServeMux()
 	mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		r, err := toRequest(req)
@@ -61,43 +86,24 @@ func (r *runnerHTTP) Run(ctx context.Context, logger *slog.Logger, h Handler) {
 		}
 	}))
 
-	s := &http.Server{
-		Addr:           fmt.Sprintf(":%d", port()),
-		Handler:        mux,
-		MaxHeaderBytes: mb,
-	}
-	go func() {
-		select {
-		case <-ctx.Done():
-			shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-			defer cancel()
-
-			logger.Info("shutting down HTTP server...")
-			if err := s.Shutdown(shutdownCtx); err != nil {
-				logger.Error("failed to shutdown server", "err", err)
-			}
-		}
-	}()
+	return mux
+}
 
-	logger.Info("serving HTTP server on port " + strconv.Itoa(port()))
-	if err := s.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-		logger.Error("unexpected shutdown of server", "err", err)
-	}
+type httpPayload struct {
+	Body        json.RawMessage `json:"body,omitempty"`
+	Context     json.RawMessage `json:"context,omitempty"`
+	AccessToken string          `json:"access_token"`
+	Method      string          `json:"method"`
+	Params      struct {
+		Header http.Header `json:"header"`
+		Query  url.Values  `json:"query"`
+	} `json:"params"`
+	URL     string `json:"url"`
+	TraceID string `json:"trace_id"`
 }
 
 func toRequest(req *http.Request) (Request, error) {
-	var r struct {
-		Body        json.RawMessage `json:"body"`
-		Context     json.RawMessage `json:"context"`
-		AccessToken string          `json:"access_token"`
-		Method      string          `json:"method"`
-		Params      struct {
-			Header http.Header `json:"header"`
-			Query  url.Values  `json:"query"`
-		} `json:"params"`
-		URL     string `json:"url"`
-		TraceID string `json:"trace_id"`
-	}
+	var r httpPayload
 	payload, err := io.ReadAll(io.LimitReader(req.Body, 5*mb))
 	if err != nil {
 		return Request{}, fmt.Errorf("failed to read request body: %s", err)
diff --git a/runner_http_direct.go b/runner_http_direct.go
new file mode 100644
index 0000000..2eee38e
--- /dev/null
+++ b/runner_http_direct.go
@@ -0,0 +1,64 @@
+package fdk
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+)
+
+// TestRun executes the handler in order to test it through the HTTP runner.
+func TestRun[T Cfg](ctx context.Context, cfg T, newHandlerFn func(_ context.Context, cfg T) Handler) (testHandlerFn func(context.Context, Request) Response) {
+	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{}))
+
+	mux := newHTTPRunMux(ctx, logger, newHandlerFn(ctx, cfg))
+
+	return func(ctx context.Context, r Request) Response {
+		b, err := json.Marshal(httpPayload{
+			Body:        r.Body,
+			Context:     r.Context,
+			AccessToken: r.AccessToken,
+			Method:      r.Method,
+			Params: struct {
+				Header http.Header `json:"header"`
+				Query  url.Values  `json:"query"`
+			}{Header: r.Params.Header, Query: r.Params.Query},
+			URL:     r.URL,
+			TraceID: r.TraceID,
+		})
+		if err != nil {
+			return ErrResp(APIError{Code: http.StatusInternalServerError, Message: "unable to create http request body: " + err.Error()})
+		}
+
+		req, err := http.NewRequestWithContext(ctx, "POST", "/", bytes.NewReader(b))
+		if err != nil {
+			return ErrResp(APIError{Code: http.StatusInternalServerError, Message: "unable to create http request: " + err.Error()})
+		}
+
+		rec := httptest.NewRecorder()
+
+		mux.ServeHTTP(rec, req)
+
+		var respBody struct {
+			Body    json.RawMessage `json:"body,omitempty"`
+			Code    int             `json:"code"`
+			Errs    []APIError      `json:"errors"`
+			Headers http.Header     `json:"headers"`
+		}
+		err = json.Unmarshal(rec.Body.Bytes(), &respBody)
+		if err != nil {
+			return ErrResp(APIError{Code: http.StatusInternalServerError, Message: "unable to unmarshal response body: " + err.Error()})
+		}
+
+		return Response{
+			Body:   respBody.Body,
+			Code:   respBody.Code,
+			Errors: respBody.Errs,
+			Header: respBody.Headers,
+		}
+	}
+}