package main import ( "context" "io" "net/http" "net/http/httptest" "strings" "sync/atomic" "testing" "time" ) func withJobServer(t *testing.T, h http.HandlerFunc) { t.Helper() srv := httptest.NewServer(h) jobHost.Lock() oldBase, oldTok := jobHost.base, jobHost.token jobHost.base, jobHost.token = srv.URL, "t" jobHost.Unlock() oldWait := putRetryWait putRetryWait = []time.Duration{time.Millisecond, time.Millisecond, time.Millisecond, time.Millisecond} t.Cleanup(func() { srv.Close() jobHost.Lock() jobHost.base, jobHost.token = oldBase, oldTok jobHost.Unlock() putRetryWait = oldWait }) } func openText(s string) func() (io.ReadCloser, int64, error) { return func() (io.ReadCloser, int64, error) { return io.NopCloser(strings.NewReader(s)), int64(len(s)), nil } } func TestPutOutputRetriesAfterAProxyError(t *testing.T) { var calls, stored atomic.Int32 withJobServer(t, func(w http.ResponseWriter, r *http.Request) { if calls.Add(1) <= 2 { http.Error(w, "bad gateway", 502) return } b, _ := io.ReadAll(r.Body) if string(b) == "video" { stored.Add(1) } w.Write([]byte("ok")) }) n, err := putOutputWithRetry(context.Background(), "j", "out.mp4", openText("video")) if err != nil || n != 5 || calls.Load() != 3 || stored.Load() != 1 { t.Fatalf("n=%d err=%v calls=%d stored=%d", n, err, calls.Load(), stored.Load()) } } func TestPutOutputDoesNotRetryARefusal(t *testing.T) { var calls atomic.Int32 withJobServer(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); http.Error(w, "no such job", 404) }) _, err := putOutputWithRetry(context.Background(), "j", "out.mp4", openText("video")) if err == nil || !strings.Contains(err.Error(), "404") || calls.Load() != 1 { t.Fatalf("err=%v calls=%d", err, calls.Load()) } } func TestPutOutputGivesUpAfterTheLastTry(t *testing.T) { var calls atomic.Int32 withJobServer(t, func(w http.ResponseWriter, r *http.Request) { calls.Add(1); http.Error(w, "x", 502) }) _, err := putOutputWithRetry(context.Background(), "j", "out.mp4", openText("video")) if err == nil || !strings.Contains(err.Error(), "502") || calls.Load() != 5 { t.Fatalf("err=%v calls=%d", err, calls.Load()) } }