Updated functionality to support context cancellation + tests

davesavic · Jan 7, 2024 · acd3579 · acd3579
1 parent 40839e0
commit acd3579
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 62 deletions.
diff --git a/client.go b/client.go
@@ -58,13 +58,20 @@ func (c *Client) Do(req *http.Request) (*http.Response, error) {
 	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
 		resp, err = c.HttpClient.Do(req)
 
+		if req.Context().Err() != nil {
+			return nil, fmt.Errorf("request context error: %w", req.Context().Err())
+		}
+
 		if c.ShouldRetryFunc != nil && !c.ShouldRetryFunc(req, resp, err) {
 			break
 		}
 
 		if attempt < c.MaxRetries {
-			// Exponential backoff only if we're going to retry.
-			time.Sleep(time.Duration(attempt) * time.Second)
+			select {
+			case <-time.After(time.Duration(attempt) * time.Second):
+			case <-req.Context().Done():
+				return nil, req.Context().Err()
+			}
 		}
 	}
 

diff --git a/client_test.go b/client_test.go
@@ -1,8 +1,10 @@
 package clink_test
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -521,3 +523,64 @@ func TestUnsuccessfulRetries(t *testing.T) {
 		t.Errorf("expected %d retries (total requests: %d), but got %d", retryCount, retryCount+1, requestCount)
 	}
 }
+
+func TestContextCancellationDuringRetries(t *testing.T) {
+	var requestCount int
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestCount++
+		w.WriteHeader(http.StatusInternalServerError) // Always return an error to trigger retries
+	}))
+	defer server.Close()
+
+	client := clink.NewClient(
+		clink.WithRetries(3, func(request *http.Request, response *http.Response, err error) bool {
+			// Always return true to retry
+			return true
+		}),
+		clink.WithClient(server.Client()),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
+	if err != nil {
+		t.Fatalf("failed to create request: %v", err)
+	}
+
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		cancel()
+	}()
+
+	_, err = client.Do(req)
+
+	if requestCount > 2 {
+		t.Errorf("expected at most 2 requests due to context cancellation, but got %d", requestCount)
+	}
+
+	if err == nil || !errors.Is(err, context.Canceled) {
+		t.Errorf("expected context cancellation error, but got: %v", err)
+	}
+}
+
+func TestRequestWithCanceledContext(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(2 * time.Second) // Simulate a delay in the response
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	client := clink.NewClient(clink.WithClient(server.Client()))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
+	if err != nil {
+		t.Fatalf("failed to create request: %v", err)
+	}
+	cancel() // Cancel the context immediately
+
+	_, err = client.Do(req)
+
+	if err == nil || !errors.Is(err, context.Canceled) {
+		t.Errorf("expected context cancellation error, but got: %v", err)
+	}
+}
diff --git a/coverage.out b/coverage.out
@@ -1,65 +1,69 @@
 mode: count
-github.com/davesavic/clink/client.go:24.40,27.27 2 24
-github.com/davesavic/clink/client.go:27.27,29.3 1 28
-github.com/davesavic/clink/client.go:31.2,31.10 1 24
-github.com/davesavic/clink/client.go:34.30,39.2 1 24
-github.com/davesavic/clink/client.go:44.64,45.36 1 16
+github.com/davesavic/clink/client.go:24.40,27.27 2 26
+github.com/davesavic/clink/client.go:27.27,29.3 1 31
+github.com/davesavic/clink/client.go:31.2,31.10 1 26
+github.com/davesavic/clink/client.go:34.30,39.2 1 26
+github.com/davesavic/clink/client.go:44.64,45.36 1 18
 github.com/davesavic/clink/client.go:45.36,47.3 1 2
-github.com/davesavic/clink/client.go:49.2,49.26 1 16
+github.com/davesavic/clink/client.go:49.2,49.26 1 18
 github.com/davesavic/clink/client.go:49.26,50.59 1 2
 github.com/davesavic/clink/client.go:50.59,52.4 1 0
-github.com/davesavic/clink/client.go:55.2,58.55 3 16
-github.com/davesavic/clink/client.go:58.55,61.69 2 19
-github.com/davesavic/clink/client.go:61.69,62.9 1 1
-github.com/davesavic/clink/client.go:65.3,65.29 1 18
-github.com/davesavic/clink/client.go:65.29,68.4 1 3
-github.com/davesavic/clink/client.go:71.2,71.16 1 16
-github.com/davesavic/clink/client.go:71.16,73.3 1 0
-github.com/davesavic/clink/client.go:75.2,75.18 1 16
-github.com/davesavic/clink/client.go:79.59,81.16 2 1
-github.com/davesavic/clink/client.go:81.16,83.3 1 0
-github.com/davesavic/clink/client.go:84.2,84.18 1 1
-github.com/davesavic/clink/client.go:88.62,90.16 2 1
-github.com/davesavic/clink/client.go:90.16,92.3 1 0
-github.com/davesavic/clink/client.go:93.2,93.18 1 1
-github.com/davesavic/clink/client.go:97.58,99.16 2 1
-github.com/davesavic/clink/client.go:99.16,101.3 1 0
-github.com/davesavic/clink/client.go:102.2,102.18 1 1
-github.com/davesavic/clink/client.go:106.75,108.16 2 1
-github.com/davesavic/clink/client.go:108.16,110.3 1 0
-github.com/davesavic/clink/client.go:111.2,111.18 1 1
-github.com/davesavic/clink/client.go:115.74,117.16 2 1
-github.com/davesavic/clink/client.go:117.16,119.3 1 0
-github.com/davesavic/clink/client.go:120.2,120.18 1 1
-github.com/davesavic/clink/client.go:124.76,126.16 2 1
-github.com/davesavic/clink/client.go:126.16,128.3 1 0
-github.com/davesavic/clink/client.go:129.2,129.18 1 1
-github.com/davesavic/clink/client.go:133.61,135.16 2 1
-github.com/davesavic/clink/client.go:135.16,137.3 1 0
-github.com/davesavic/clink/client.go:138.2,138.18 1 1
-github.com/davesavic/clink/client.go:144.45,145.25 1 16
-github.com/davesavic/clink/client.go:145.25,147.3 1 16
-github.com/davesavic/clink/client.go:151.43,152.25 1 2
-github.com/davesavic/clink/client.go:152.25,154.3 1 2
-github.com/davesavic/clink/client.go:158.52,159.25 1 2
-github.com/davesavic/clink/client.go:159.25,160.35 1 2
-github.com/davesavic/clink/client.go:160.35,162.4 1 2
-github.com/davesavic/clink/client.go:167.36,168.25 1 2
-github.com/davesavic/clink/client.go:168.25,171.3 2 2
-github.com/davesavic/clink/client.go:175.54,176.25 1 1
-github.com/davesavic/clink/client.go:176.25,180.3 3 1
-github.com/davesavic/clink/client.go:184.42,185.25 1 1
-github.com/davesavic/clink/client.go:185.25,187.3 1 1
-github.com/davesavic/clink/client.go:191.38,192.25 1 1
+github.com/davesavic/clink/client.go:55.2,58.55 3 18
+github.com/davesavic/clink/client.go:58.55,61.33 2 22
+github.com/davesavic/clink/client.go:61.33,63.4 1 1
+github.com/davesavic/clink/client.go:65.3,65.69 1 21
+github.com/davesavic/clink/client.go:65.69,66.9 1 1
+github.com/davesavic/clink/client.go:69.3,69.29 1 20
+github.com/davesavic/clink/client.go:69.29,70.11 1 5
+github.com/davesavic/clink/client.go:71.60,71.60 0 4
+github.com/davesavic/clink/client.go:72.32,73.36 1 1
+github.com/davesavic/clink/client.go:78.2,78.16 1 16
+github.com/davesavic/clink/client.go:78.16,80.3 1 0
+github.com/davesavic/clink/client.go:82.2,82.18 1 16
+github.com/davesavic/clink/client.go:86.59,88.16 2 1
+github.com/davesavic/clink/client.go:88.16,90.3 1 0
+github.com/davesavic/clink/client.go:91.2,91.18 1 1
+github.com/davesavic/clink/client.go:95.62,97.16 2 1
+github.com/davesavic/clink/client.go:97.16,99.3 1 0
+github.com/davesavic/clink/client.go:100.2,100.18 1 1
+github.com/davesavic/clink/client.go:104.58,106.16 2 1
+github.com/davesavic/clink/client.go:106.16,108.3 1 0
+github.com/davesavic/clink/client.go:109.2,109.18 1 1
+github.com/davesavic/clink/client.go:113.75,115.16 2 1
+github.com/davesavic/clink/client.go:115.16,117.3 1 0
+github.com/davesavic/clink/client.go:118.2,118.18 1 1
+github.com/davesavic/clink/client.go:122.74,124.16 2 1
+github.com/davesavic/clink/client.go:124.16,126.3 1 0
+github.com/davesavic/clink/client.go:127.2,127.18 1 1
+github.com/davesavic/clink/client.go:131.76,133.16 2 1
+github.com/davesavic/clink/client.go:133.16,135.3 1 0
+github.com/davesavic/clink/client.go:136.2,136.18 1 1
+github.com/davesavic/clink/client.go:140.61,142.16 2 1
+github.com/davesavic/clink/client.go:142.16,144.3 1 0
+github.com/davesavic/clink/client.go:145.2,145.18 1 1
+github.com/davesavic/clink/client.go:151.45,152.25 1 18
+github.com/davesavic/clink/client.go:152.25,154.3 1 18
+github.com/davesavic/clink/client.go:158.43,159.25 1 2
+github.com/davesavic/clink/client.go:159.25,161.3 1 2
+github.com/davesavic/clink/client.go:165.52,166.25 1 2
+github.com/davesavic/clink/client.go:166.25,167.35 1 2
+github.com/davesavic/clink/client.go:167.35,169.4 1 2
+github.com/davesavic/clink/client.go:174.36,175.25 1 2
+github.com/davesavic/clink/client.go:175.25,178.3 2 2
+github.com/davesavic/clink/client.go:182.54,183.25 1 1
+github.com/davesavic/clink/client.go:183.25,187.3 3 1
+github.com/davesavic/clink/client.go:191.42,192.25 1 1
 github.com/davesavic/clink/client.go:192.25,194.3 1 1
-github.com/davesavic/clink/client.go:198.95,199.25 1 3
-github.com/davesavic/clink/client.go:199.25,202.3 2 3
-github.com/davesavic/clink/client.go:206.70,207.21 1 7
-github.com/davesavic/clink/client.go:207.21,209.3 1 1
-github.com/davesavic/clink/client.go:211.2,211.26 1 6
-github.com/davesavic/clink/client.go:211.26,213.3 1 1
-github.com/davesavic/clink/client.go:215.2,215.33 1 5
-github.com/davesavic/clink/client.go:215.33,217.3 1 5
-github.com/davesavic/clink/client.go:219.2,219.70 1 5
-github.com/davesavic/clink/client.go:219.70,221.3 1 1
-github.com/davesavic/clink/client.go:223.2,223.12 1 4
+github.com/davesavic/clink/client.go:198.38,199.25 1 1
+github.com/davesavic/clink/client.go:199.25,201.3 1 1
+github.com/davesavic/clink/client.go:205.95,206.25 1 4
+github.com/davesavic/clink/client.go:206.25,209.3 2 4
+github.com/davesavic/clink/client.go:213.70,214.21 1 7
+github.com/davesavic/clink/client.go:214.21,216.3 1 1
+github.com/davesavic/clink/client.go:218.2,218.26 1 6
+github.com/davesavic/clink/client.go:218.26,220.3 1 1
+github.com/davesavic/clink/client.go:222.2,222.33 1 5
+github.com/davesavic/clink/client.go:222.33,224.3 1 5
+github.com/davesavic/clink/client.go:226.2,226.70 1 5
+github.com/davesavic/clink/client.go:226.70,228.3 1 1
+github.com/davesavic/clink/client.go:230.2,230.12 1 4