go-omni/httpc/client_test.go

package httpc

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/logx"
)

type trackedBody struct {
	io.Reader
	closed bool
}

func (b *trackedBody) Close() error { b.closed = true; return nil }

func TestRejectedRequestClosesBody(t *testing.T) {
	for _, closed := range []bool{false, true} {
		client, _ := New(Config{DisableLogging: true})
		body := &trackedBody{Reader: strings.NewReader("body")}
		request, _ := http.NewRequest("POST", "http://unused.invalid", body)
		if closed {
			client.Close()
		} else {
			request.URL = nil
		}
		_, err := client.Do(request)
		client.Close()
		if err == nil || !body.closed {
			t.Fatalf("error=%v bodyClosed=%v", err, body.closed)
		}
	}
}

func TestRetryCorrelationPoolAndWriteSafety(t *testing.T) {
	var calls, connections atomic.Int32
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := calls.Add(1)
		if r.Header.Get("X-Request-ID") != "request-1" {
			t.Error("missing correlation")
		}
		if n == 1 || r.Method == "POST" {
			w.WriteHeader(503)
			io.WriteString(w, "busy")
			return
		}
		io.WriteString(w, `{"ok":true}`)
	}))
	server.Config.ConnState = func(_ net.Conn, s http.ConnState) {
		if s == http.StateNew {
			connections.Add(1)
		}
	}
	server.Start()
	defer server.Close()
	client, err := New(Config{MaxRetries: 2, Backoff: time.Millisecond, DisableLogging: true})
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	ctx := logx.WithRequestID(context.Background(), "request-1")
	var output struct {
		OK bool `json:"ok"`
	}
	if err := client.JSON(ctx, "GET", server.URL, nil, &output); err != nil || !output.OK {
		t.Fatal(err)
	}
	if calls.Load() != 2 || connections.Load() != 1 {
		t.Fatalf("calls=%d connections=%d", calls.Load(), connections.Load())
	}
	if err := client.JSON(ctx, "POST", server.URL, map[string]string{"name": "value"}, nil); err == nil {
		t.Fatal("expected 503")
	}
	if calls.Load() != 3 {
		t.Fatal("POST was retried")
	}
	client.Close()
	if err := client.JSON(ctx, "GET", server.URL, nil, nil); !errors.Is(err, ErrClosed) {
		t.Fatal(err)
	}
}

func TestTotalDeadlineIncludesBackoffAndBody(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/retry" {
			w.Header().Set("Retry-After", "1")
			w.WriteHeader(503)
			return
		}
		w.WriteHeader(200)
		w.(http.Flusher).Flush()
		<-r.Context().Done()
	}))
	defer server.Close()
	client, _ := New(Config{Timeout: 40 * time.Millisecond, MaxRetries: 3, DisableLogging: true})
	defer client.Close()
	started := time.Now()
	err := client.JSON(context.Background(), "GET", server.URL+"/retry", nil, nil)
	if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > time.Second {
		t.Fatal(err, time.Since(started))
	}
	req, _ := http.NewRequest("GET", server.URL+"/body", nil)
	response, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	if _, err := io.ReadAll(response.Body); err == nil {
		t.Fatal("body deadline was canceled too early")
	}
}

func TestLimitsRedirectAndNonReplayableBody(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		switch r.URL.Path {
		case "/redirect":
			http.Redirect(w, r, "/large", 302)
		case "/large":
			io.WriteString(w, strings.Repeat("x", 17))
		default:
			w.WriteHeader(503)
		}
	}))
	defer server.Close()
	client, _ := New(Config{MaxRetries: 2, MaxResponseBytes: 16, DisableLogging: true})
	defer client.Close()
	if err := client.JSON(context.Background(), "GET", server.URL+"/large", nil, nil); !errors.Is(err, ErrResponseTooLarge) {
		t.Fatal(err)
	}
	var status *StatusError
	if err := client.JSON(context.Background(), "GET", server.URL+"/redirect", nil, nil); !errors.As(err, &status) || status.StatusCode != 302 {
		t.Fatal(err)
	}
	if calls.Load() != 2 {
		t.Fatal("followed redirect")
	}
	req, _ := http.NewRequest("GET", server.URL, io.NopCloser(strings.NewReader("body")))
	res, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()
	if calls.Load() != 3 {
		t.Fatal("non-replayable body was retried")
	}
}