go-omni/httpc/status_test.go

package httpc

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

func TestJSONSuccessStatuses(t *testing.T) {
	for _, status := range []int{200, 201, 204, 264, 299, 400} {
		t.Run(strconv.Itoa(status), func(t *testing.T) {
			var calls atomic.Int32
			body := &trackedBody{Reader: strings.NewReader(`{"value":2}`)}
			c := testClient(t, Config{MaxRetries: 5, Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
				calls.Add(1)
				return &http.Response{StatusCode: status, Body: body, ContentLength: 11, Header: make(http.Header), Request: r}, nil
			})})
			out := struct {
				Value int `json:"value"`
			}{Value: 1}
			err := c.JSON(context.Background(), "GET", "http://example.invalid", nil, &out, WithSuccessStatuses(200))
			if status == 200 {
				if err != nil || out.Value != 2 {
					t.Fatal(err, out)
				}
			} else {
				var rejected *StatusError
				if !errors.As(err, &rejected) || rejected.StatusCode != status || out.Value != 1 {
					t.Fatal(err, out)
				}
				if strings.Contains(err.Error(), "value") {
					t.Fatal("status error exposed body")
				}
			}
			if !body.closed || calls.Load() != 1 {
				t.Fatal("not closed or unexpected retry", calls.Load())
			}
			// Unconfigured calls continue to accept every 2xx.
			c2 := testClient(t, Config{Transport: benchmarkTransport{`{}`, status}})
			err = c2.JSON(context.Background(), "GET", "http://example.invalid", nil, nil)
			if (err == nil) != (status >= 200 && status < 300) {
				t.Fatal(err)
			}
			if status >= 200 && status < 300 {
				if err := c2.JSON(context.Background(), "GET", "http://example.invalid", nil, nil, WithSuccessStatuses(status, status)); err != nil {
					t.Fatal(err)
				}
			}
		})
	}
	var calls atomic.Int32
	c := testClient(t, Config{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
		calls.Add(1)
		return benchmarkTransport{`{}`, 200}.RoundTrip(r)
	})})
	for _, option := range []RequestOption{WithSuccessStatuses(), WithSuccessStatuses(199), WithSuccessStatuses(300), WithSuccessStatuses(-1), WithSuccessStatuses(200, 500)} {
		if err := c.JSON(context.Background(), "GET", "http://example.invalid", nil, nil, option); err == nil {
			t.Fatal("invalid policy accepted")
		}
	}
	if calls.Load() != 0 {
		t.Fatal("invalid policy reached transport")
	}
}
func TestJSONSuccessStatusConcurrentOptions(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n, _ := strconv.Atoi(r.URL.Query().Get("status"))
		w.WriteHeader(n)
		if n != 204 {
			io.WriteString(w, `{}`)
		}
	}))
	defer srv.Close()
	c := testClient(t, Config{})
	statuses := []int{200, 201}
	option := WithSuccessStatuses(statuses...)
	statuses[0] = 204
	var wg sync.WaitGroup
	for i := 0; i < 30; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			status := 200 + i%2
			err := c.JSON(context.Background(), "GET", fmt.Sprintf("%s?status=%d", srv.URL, status), nil, nil, option)
			if err != nil {
				t.Error(err)
			}
			err = c.JSON(context.Background(), "GET", srv.URL+"?status=204", nil, nil, option)
			var rejected *StatusError
			if !errors.As(err, &rejected) || rejected.StatusCode != 204 {
				t.Error(err)
			}
		}(i)
	}
	wg.Wait()
	// Last explicit policy wins; another call's policy must not leak.
	if err := c.JSON(context.Background(), "GET", srv.URL+"?status=204", nil, nil, option, WithSuccessStatuses(204)); err != nil {
		t.Fatal(err)
	}
	if err := c.JSON(context.Background(), "GET", srv.URL+"?status=204", nil, nil); err != nil {
		t.Fatal(err)
	}
}
func TestJSONRejectedSuccessBodyRecovery(t *testing.T) {
	var connections atomic.Int32
	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(201)
		if r.URL.Path == "/stream" {
			w.(http.Flusher).Flush()
			<-r.Context().Done()
			return
		}
		io.WriteString(w, "private response")
	}))
	srv.Config.ConnState = func(_ net.Conn, s http.ConnState) {
		if s == http.StateNew {
			connections.Add(1)
		}
	}
	srv.Start()
	defer srv.Close()
	c := testClient(t, Config{Timeout: time.Second})
	for i := 0; i < 5; i++ {
		err := c.JSON(context.Background(), "GET", srv.URL, nil, nil, WithSuccessStatuses(200))
		var status *StatusError
		if !errors.As(err, &status) || status.StatusCode != 201 {
			t.Fatal(err)
		}
	}
	if connections.Load() != 1 {
		t.Fatal("small rejected bodies lost keep-alive", connections.Load())
	}
	start := time.Now()
	err := c.JSON(context.Background(), "GET", srv.URL+"/stream", nil, nil, WithSuccessStatuses(200))
	var status *StatusError
	if !errors.As(err, &status) || status.StatusCode != 201 || time.Since(start) > 500*time.Millisecond {
		t.Fatal("waited for rejected unbounded body", err)
	}
	if err := c.JSON(context.Background(), "GET", srv.URL, nil, nil, WithSuccessStatuses(201), WithMaxResponseBytes(4)); !errors.Is(err, ErrResponseTooLarge) {
		t.Fatal("accepted body bypassed size limit", err)
	}
}
func BenchmarkJSONSuccessStatus(b *testing.B) {
	for _, restrict := range []bool{false, true} {
		b.Run(strconv.FormatBool(restrict), func(b *testing.B) {
			c, _ := New(Config{Transport: benchmarkTransport{`{"id":1}`, 200}, DisableLogging: true})
			defer c.Close()
			var options []RequestOption
			if restrict {
				options = []RequestOption{WithSuccessStatuses(200)}
			}
			var out struct {
				ID int `json:"id"`
			}
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				if err := c.JSON(context.Background(), "GET", "http://example.invalid", nil, &out, options...); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}