Update NewClient to use functional options pattern (#3)

* Update NewClient to use functional options pattern * Reorder declarations * Add WithHTTPClient * Use http.DefaultClient * Use replicate.NewClient in tests
replicate · Jul 14, 2023 · cb77e43 · cb77e43
1 parent 11a47ef
commit cb77e43
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 83 deletions.
diff --git a/replicate/client.go b/replicate/client.go
@@ -10,38 +10,55 @@ import (
 	"strings"
 )
 
-// Client represents a Replicate API client.
+// Client is a client for the Replicate API.
 type Client struct {
 	Auth       string
 	UserAgent  *string
 	BaseURL    string
 	HTTPClient *http.Client
 }
 
+// ClientOption is a function that modifies a Client.
+type ClientOption func(*Client)
+
 // NewClient creates a new Replicate API client.
-func NewClient(auth string) *Client {
-	return NewClientWithOptions(auth, nil, nil)
-}
+func NewClient(auth string, options ...ClientOption) *Client {
+	defaultUserAgent := "replicate-go"
+	defaultBaseURL := "https://api.replicate.com/v1"
+	defaultClient := http.DefaultClient
 
-// NewClientWithOptions creates a new Replicate API client with a custom user agent and base URL.
-func NewClientWithOptions(auth string, userAgent *string, baseURL *string) *Client {
-	client := &http.Client{}
+	c := &Client{
+		Auth:       auth,
+		UserAgent:  &defaultUserAgent,
+		BaseURL:    defaultBaseURL,
+		HTTPClient: defaultClient,
+	}
 
-	if userAgent == nil {
-		defaultUserAgent := "replicate-go"
-		userAgent = &defaultUserAgent
+	for _, option := range options {
+		option(c)
 	}
 
-	if baseURL == nil {
-		defaultBaseURL := "https://api.replicate.com/v1"
-		baseURL = &defaultBaseURL
+	return c
+}
+
+// WithUserAgent sets the User-Agent header on requests made by the client.
+func WithUserAgent(userAgent string) ClientOption {
+	return func(c *Client) {
+		c.UserAgent = &userAgent
 	}
+}
 
-	return &Client{
-		Auth:       auth,
-		UserAgent:  userAgent,
-		BaseURL:    *baseURL,
-		HTTPClient: client,
+// WithBaseURL sets the base URL for the client.
+func WithBaseURL(baseURL string) ClientOption {
+	return func(c *Client) {
+		c.BaseURL = baseURL
+	}
+}
+
+// WithHTTPClient sets the HTTP client used by the client.
+func WithHTTPClient(httpClient *http.Client) ClientOption {
+	return func(c *Client) {
+		c.HTTPClient = httpClient
 	}
 }
 

diff --git a/replicate/replicate_test.go b/replicate/replicate_test.go
@@ -53,11 +53,9 @@ func TestListCollections(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	initialPage, err := client.ListCollections(context.Background())
 	if err != nil {
@@ -107,11 +105,9 @@ func TestGetCollection(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -149,11 +145,9 @@ func TestGetModel(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	model, err := client.GetModel(context.Background(), "replicate", "hello-world")
 	assert.NoError(t, err)
@@ -180,11 +174,9 @@ func TestListModelVersions(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	versionsPage, err := client.ListModelVersions(context.Background(), "replicate", "hello-world")
 	assert.NoError(t, err)
@@ -211,11 +203,9 @@ func TestGetModelVersion(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	version, err := client.GetModelVersion(context.Background(), "replicate", "hello-world", "version1")
 	assert.NoError(t, err)
@@ -266,11 +256,9 @@ func TestCreatePrediction(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	input := replicate.PredictionInput{"text": "Alice"}
 	webhook := replicate.Webhook{
@@ -329,11 +317,9 @@ func TestListPredictions(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	initialPage, err := client.ListPredictions(context.Background())
 	if err != nil {
@@ -382,11 +368,9 @@ func TestGetPrediction(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -430,11 +414,9 @@ func TestWait(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	prediction := &replicate.Prediction{
 		ID:        "ufawqhfynnddngldkgtslldrkq",
@@ -475,11 +457,9 @@ func TestCreateTraining(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -517,11 +497,9 @@ func TestGetTraining(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -552,11 +530,9 @@ func TestCancelTraining(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -589,11 +565,9 @@ func TestListTrainings(t *testing.T) {
 	}))
 	defer mockServer.Close()
 
-	client := &replicate.Client{
-		BaseURL:    mockServer.URL,
-		Auth:       "test-token",
-		HTTPClient: http.DefaultClient,
-	}
+	client := replicate.NewClient("test-token",
+		replicate.WithBaseURL(mockServer.URL),
+	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()