feat: option-to-pass-optional-resty-or-http-client

Author: wood

client.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"net/http"
+	"reflect"
+	"runtime"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -34,6 +37,7 @@ type Client struct {
 	defaultFlagHandler func(string) (Flag, error)
 
 	client         *resty.Client
+	httpClient     *http.Client
 	ctxLocalEval   context.Context
 	ctxAnalytics   context.Context
 	log            *slog.Logger
@@ -52,12 +56,42 @@ func GetEvaluationContextFromCtx(ctx context.Context) (ec EvaluationContext, ok
 	return ec, ok
 }
 
+func getOptionQualifiedName(opt Option) string {
+	return runtime.FuncForPC(reflect.ValueOf(opt).Pointer()).Name()
+}
+
+func isClientOption(name string) bool {
+	return strings.Contains(name, OptionWithHTTPClient) || strings.Contains(name, OptionWithRestyClient)
+}
+
 // NewClient creates instance of Client with given configuration.
 func NewClient(apiKey string, options ...Option) *Client {
 	c := &Client{
 		apiKey: apiKey,
 		config: defaultConfig(),
-		client: resty.New(),
+	}
+
+	for _, opt := range options {
+		name := getOptionQualifiedName(opt)
+		if isClientOption(name) {
+			opt(c)
+		}
+	}
+
+	if c.client == nil {
+		if c.httpClient != nil {
+			c.client = resty.NewWithClient(c.httpClient)
+		} else {
+			c.client = resty.New()
+		}
+	}
+
+	for _, opt := range options {
+		name := getOptionQualifiedName(opt)
+		if isClientOption(name) {
+			continue
+		}
+		opt(c)
 	}
 
 	c.client.SetHeaders(map[string]string{
@@ -72,6 +106,16 @@ func NewClient(apiKey string, options ...Option) *Client {
 			opt(c)
 		}
 	}
+
+	// If a resty custom client has been provided, client is already set - otherwise we use a custom http client or default to a resty
+	if c.client == nil {
+		if c.httpClient != nil {
+			c.client = resty.NewWithClient(c.httpClient)
+		} else {
+			c.client = resty.New()
+		}
+	}
+
 	c.client = c.client.
 		SetLogger(newSlogToRestyAdapter(c.log)).
 		OnBeforeRequest(newRestyLogRequestMiddleware(c.log)).

client_test.go

@@ -15,6 +15,7 @@ import (
 
 	flagsmith "github.com/Flagsmith/flagsmith-go-client/v4"
 	"github.com/Flagsmith/flagsmith-go-client/v4/fixtures"
+	"github.com/go-resty/resty/v2"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -1019,3 +1020,140 @@ type writerFunc func(p []byte) (n int, err error)
 func (f writerFunc) Write(p []byte) (n int, err error) {
 	return f(p)
 }
+
+// Helper functions to implement a header interceptor
+func roundTripperWithHeader(key, value string) http.RoundTripper {
+	return &injectHeaderTransport{key: key, value: value}
+}
+
+type injectHeaderTransport struct {
+	key   string
+	value string
+}
+
+func (t *injectHeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set(t.key, t.value)
+	return http.DefaultTransport.RoundTrip(req)
+}
+
+func TestCustomHTTPClientIsUsed(t *testing.T) {
+	ctx := context.Background()
+
+	hasCustomHeader := false
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		assert.Equal(t, "/api/v1/flags/", req.URL.Path)
+		assert.Equal(t, fixtures.EnvironmentAPIKey, req.Header.Get("x-Environment-Key"))
+		if req.Header.Get("X-Test-Client") == "http" {
+			hasCustomHeader = true
+		}
+		rw.Header().Set("Content-Type", "application/json")
+		rw.WriteHeader(http.StatusOK)
+		_, err := io.WriteString(rw, fixtures.FlagsJson)
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	customClient := &http.Client{
+		Transport: roundTripperWithHeader("X-Test-Client", "http"),
+	}
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithHTTPClient(customClient),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	flags, err := client.GetFlags(ctx, nil)
+	assert.Equal(t, 1, len(flags.AllFlags()))
+	assert.NoError(t, err)
+	assert.True(t, hasCustomHeader, "Expected http header")
+	flag, err := flags.GetFlag(fixtures.Feature1Name)
+	assert.NoError(t, err)
+	assert.Equal(t, fixtures.Feature1Value, flag.Value)
+}
+
+func TestCustomRestyClientIsUsed(t *testing.T) {
+	ctx := context.Background()
+
+	hasCustomHeader := false
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if req.Header.Get("X-Custom-Test-Header") == "resty" {
+			hasCustomHeader = true
+		}
+		rw.Header().Set("Content-Type", "application/json")
+		rw.WriteHeader(http.StatusOK)
+		_, err := io.WriteString(rw, fixtures.FlagsJson)
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	restyClient := resty.New().
+		SetHeader("X-Custom-Test-Header", "resty")
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithRestyClient(restyClient),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	flags, err := client.GetFlags(ctx, nil)
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(flags.AllFlags()))
+	assert.True(t, hasCustomHeader, "Expected custom resty header")
+}
+
+func TestRestyClientOverridesHTTPClient(t *testing.T) {
+	ctx := context.Background()
+
+	var customHeader string
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		customHeader = req.Header.Get("X-Test-Client")
+		rw.Header().Set("Content-Type", "application/json")
+		rw.WriteHeader(http.StatusOK)
+		_, err := io.WriteString(rw, fixtures.FlagsJson)
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	httpClient := &http.Client{
+		Transport: roundTripperWithHeader("X-Test-Client", "http"),
+	}
+
+	restyClient := resty.New().
+		SetHeader("X-Test-Client", "resty")
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithHTTPClient(httpClient),
+		flagsmith.WithRestyClient(restyClient),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	flags, err := client.GetFlags(ctx, nil)
+
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(flags.AllFlags()))
+	assert.Equal(t, "resty", customHeader, "Expected resty header")
+}
+
+func TestDefaultRestyClientIsUsed(t *testing.T) {
+	ctx := context.Background()
+
+	serverCalled := false
+
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		serverCalled = true
+
+		assert.Equal(t, "/api/v1/flags/", req.URL.Path)
+		assert.Equal(t, fixtures.EnvironmentAPIKey, req.Header.Get("x-Environment-Key"))
+
+		rw.Header().Set("Content-Type", "application/json")
+		rw.WriteHeader(http.StatusOK)
+		_, err := io.WriteString(rw, fixtures.FlagsJson)
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey,
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+
+	flags, err := client.GetFlags(ctx, nil)
+
+	assert.NoError(t, err)
+	assert.True(t, serverCalled, "Expected server to be")
+	assert.Equal(t, 1, len(flags.AllFlags()))
+}

options.go

@@ -2,10 +2,18 @@ package flagsmith
 
 import (
 	"context"
+	"net/http"
 	"strings"
 	"time"
 
 	"log/slog"
+
+	"github.com/go-resty/resty/v2"
+)
+
+const (
+	OptionWithHTTPClient  = "WithHTTPClient"
+	OptionWithRestyClient = "WithRestyClient"
 )
 
 type Option func(c *Client)
@@ -165,3 +173,19 @@ func WithPolling() Option {
 		c.config.polling = true
 	}
 }
+
+func WithHTTPClient(httpClient *http.Client) Option {
+	return func(c *Client) {
+		if httpClient != nil {
+			c.httpClient = httpClient
+		}
+	}
+}
+
+func WithRestyClient(restyClient *resty.Client) Option {
+	return func(c *Client) {
+		if restyClient != nil {
+			c.client = restyClient
+		}
+	}
+}