sub2api/backend/internal/repository/account_traffic_http_test.go

package repository

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

	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/stretchr/testify/require"
)

func TestAccountTrafficHTTPEnforcesBudgetBeforeNetwork(t *testing.T) {
	_, cache, plan := trafficFixture(t)
	plan.Policy.RPM = 1
	plan.Policy.Burst = 1
	control := service.NewAccountTrafficService(cache)
	upstream := NewControlledHTTPUpstream(nil, control)
	a := &service.Account{ID: 42, Concurrency: 8, Extra: map[string]any{service.AccountTrafficPolicyKey: plan.Policy}}
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		io.WriteString(w, "data: hello\n\n")
	}))
	defer server.Close()
	req, _ := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
	resp, err := upstream.Do(service.WithAccountTrafficRequest(req, a), "", 42, 8)
	require.NoError(t, err)
	_, err = io.ReadAll(resp.Body)
	require.NoError(t, err)
	require.NoError(t, resp.Body.Close())
	req, _ = http.NewRequest("POST", server.URL, strings.NewReader("{}"))
	_, err = upstream.Do(service.WithAccountTrafficRequest(req, a), "", 42, 8)
	var limited *service.AccountTrafficLimitError
	require.True(t, errors.As(err, &limited))
	require.EqualValues(t, 1, calls.Load())
	// Turning both switches off immediately restores the normal transport.
	a.Extra = nil
	req, _ = http.NewRequest("POST", server.URL, strings.NewReader("{}"))
	resp, err = upstream.Do(service.WithAccountTrafficRequest(req, a), "", 42, 8)
	require.NoError(t, err)
	resp.Body.Close()
	require.EqualValues(t, 2, calls.Load())
}