diff --git a/internal/export/export_test.go b/internal/export/export_test.go index 68d82bd..f9ea4ef 100644 --- a/internal/export/export_test.go +++ b/internal/export/export_test.go @@ -104,6 +104,29 @@ func (f *fixture) downloads() []testserver.Event { return out } +// partSizes maps each archive of example.com that has a partial file +// ("images-1.zip") to that file's size. It reads the size through an open +// file, which is current on every platform even while a download still +// writes to it. +func (f *fixture) partSizes() map[string]int64 { + out := map[string]int64{} + paths, _ := filepath.Glob(filepath.Join(f.dir, "example.com", "*"+partSuffix)) + for _, p := range paths { + name := filepath.Base(p) + i := strings.Index(name, ".zip.") + file, err := os.Open(p) + if i < 0 || err != nil { + continue + } + st, err := file.Stat() + file.Close() + if err == nil { + out[name[:i+len(".zip")]] = st.Size() + } + } + return out +} + func TestRunDownloadsArchivesWhileTheExportRuns(t *testing.T) { f := newFixture(t) a1, a2, d1 := zipOf(t, "example.com/a.jpg"), zipOf(t, "example.com/b.jpg"), zipOf(t, "c.png") @@ -759,28 +782,39 @@ func TestRunResumesDownloadsCancelledMidway(t *testing.T) { a, b := zipOf(t, "a"), zipOf(t, "b") id := f.srv.AddExport(api.ExportCompleted, testserver.ExportDomain{Domain: "example.com", Archives: [][]byte{a, b}}) f.srv.Inject(testserver.RouteDownload, testserver.Response{StallAfter: 30}, testserver.Response{StallAfter: 40}) + // Whichever archive starts first stalls after 30 bytes, the other after + // 40. The run is cancelled once both partial files hold those bytes, + // however long the downloads take to get there. + stalled := func(parts map[string]int64) bool { + one, two := parts["images-1.zip"], parts["images-2.zip"] + return len(parts) == 2 && (one == 30 && two == 40 || one == 40 && two == 30) + } ctx, cancel := context.WithCancel(context.Background()) - stalled := 0 - f.srv.Hook = func(route testserver.Route, _ *http.Request) { - if route == testserver.RouteDownload { - f.mu.Lock() - stalled++ - if stalled == 2 { - time.AfterFunc(100*time.Millisecond, cancel) + defer cancel() + ran := make(chan struct{}) + go func() { + defer cancel() + for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline) && !stalled(f.partSizes()); { + select { + case <-ran: + return + case <-time.After(time.Millisecond): } - f.mu.Unlock() } - } + }() out := f.run(ctx, id) + close(ran) if !errors.Is(out.Err, context.Canceled) || out.Failed != 0 || out.Downloaded != 0 { t.Fatalf("cancelled run: %+v", out) } - parts, _ := filepath.Glob(filepath.Join(f.dir, "example.com", "*"+partSuffix)) - if len(parts) != 2 { + parts := f.partSizes() + if !stalled(parts) { t.Fatalf("partial files: %v", parts) } - f.srv.Hook = nil + // The fake files the cancelled downloads only once it notices the + // client left, which can be after Run returned. + f.srv.WaitIdle(t) before := len(f.downloads()) out = f.run(context.Background(), id) if out.Err != nil || out.Downloaded != 2 { @@ -789,9 +823,14 @@ func TestRunResumesDownloadsCancelledMidway(t *testing.T) { if !bytes.Equal(f.file("example.com/images-1.zip"), a) || !bytes.Equal(f.file("example.com/images-2.zip"), b) { t.Error("resumed archives differ") } - for _, d := range f.downloads()[before:] { - if d.Range != "bytes=30-" && d.Range != "bytes=40-" { - t.Errorf("resume without the right range: %+v", d) + resumed := f.downloads()[before:] + if len(resumed) != 2 { + t.Fatalf("resumed downloads: %+v", resumed) + } + for _, d := range resumed { + archive := "images-" + d.Key[strings.LastIndex(d.Key, "/")+1:] + ".zip" + if want := fmt.Sprintf("bytes=%d-", parts[archive]); d.Range != want { + t.Errorf("%s resumed with Range %q, want %q: %+v", archive, d.Range, want, d) } } } diff --git a/internal/testserver/testserver.go b/internal/testserver/testserver.go index 1294a54..ad85aff 100644 --- a/internal/testserver/testserver.go +++ b/internal/testserver/testserver.go @@ -171,6 +171,10 @@ type Server struct { violations []string inflightPut int maxPut int + // active counts the requests being handled; idle is closed once it + // drops to zero (see WaitIdle). + active int + idle chan struct{} exports map[int64]*fakeExport nextExport int64 @@ -196,14 +200,14 @@ func New(t testing.TB) *Server { api.HandleFunc("/api/folders", s.handleFolders) api.HandleFunc("/api/asset-exports", s.handleExports) api.HandleFunc("/api/asset-exports/", s.handleExport) - s.API = httptest.NewServer(api) - s.S3 = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s.API = httptest.NewServer(s.track(api)) + s.S3 = httptest.NewServer(s.track(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if strings.HasPrefix(r.URL.Path, "/exports/") { s.handleDownload(w, r) return } s.handlePut(w, r) - })) + }))) s.URL = s.API.URL + "/api" t.Cleanup(func() { s.API.Close() @@ -270,6 +274,47 @@ func (s *Server) Events() []Event { return append([]Event(nil), s.events...) } +// WaitIdle waits until no request is being handled, so Events holds every +// request whose handler has started. A request is filed when its handler +// ends, and the handler of a request the client gave up on (a stalled +// download, say) ends only once the fake notices, which can be after the +// client has moved on. The test fails if requests are still being handled +// after 10 seconds. +func (s *Server) WaitIdle(t testing.TB) { + t.Helper() + s.mu.Lock() + active, idle := s.active, s.idle + s.mu.Unlock() + if active == 0 { + return + } + select { + case <-idle: + case <-time.After(10 * time.Second): + t.Fatal("testserver: requests still being handled after 10s") + } +} + +// track counts the requests next is handling, for WaitIdle. +func (s *Server) track(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + if s.active == 0 { + s.idle = make(chan struct{}) + } + s.active++ + s.mu.Unlock() + defer func() { + s.mu.Lock() + if s.active--; s.active == 0 { + close(s.idle) + } + s.mu.Unlock() + }() + next.ServeHTTP(w, r) + }) +} + // Count returns how many requests reached a route. func (s *Server) Count(route Route) int { n := 0