cuducos.me
馃П Chunk is a download manager for slow and unstable servers
1package chunk
2
3import (
4 "archive/zip"
5 "bytes"
6 "context"
7 "io"
8 "math/rand"
9 "net/http"
10 "net/http/httptest"
11 "os"
12 "path/filepath"
13 "strings"
14 "sync/atomic"
15 "testing"
16 "time"
17)
18
19var timeout = 250 * time.Millisecond
20
21func TestDownload_HTTPFailure(t *testing.T) {
22 t.Parallel()
23 s := httptest.NewServer(http.HandlerFunc(
24 func(w http.ResponseWriter, r *http.Request) {
25 if r.Method == http.MethodHead {
26 w.Header().Add("Content-Length", "2")
27 return
28 }
29 w.WriteHeader(http.StatusBadRequest)
30 },
31 ))
32 defer s.Close()
33 d := Downloader{
34 OutputDir: t.TempDir(),
35 Timeout: timeout,
36 MaxRetries: 4,
37 ConcurrencyPerServer: 1,
38 ChunkSize: 1024,
39 WaitRetry: 1 * time.Millisecond,
40 }
41 ch := d.Download(s.URL)
42 <-ch // discard the first got (just the file size)
43 got := <-ch
44 if got.Error == nil {
45 t.Error("expected an error, but got nil")
46 }
47 if _, ok := <-ch; ok {
48 t.Error("expected channel closed, but did not get it")
49 }
50}
51
52func TestDownload_ServerTimeout(t *testing.T) {
53 t.Parallel()
54 s := httptest.NewServer(http.HandlerFunc(
55 func(w http.ResponseWriter, r *http.Request) {
56 if r.Method == http.MethodHead {
57 w.Header().Add("Content-Length", "2")
58 return
59 }
60 time.Sleep(10 * timeout) // server sleeps, causing client timeout
61 },
62 ))
63 defer s.Close()
64 d := Downloader{
65 OutputDir: t.TempDir(),
66 Timeout: timeout,
67 MaxRetries: 4,
68 ConcurrencyPerServer: 1,
69 ChunkSize: 1024,
70 WaitRetry: 1 * time.Millisecond,
71 }
72 // context with timeout to eventually terminate the download.
73 ctx, cancel := context.WithTimeout(context.Background(), 3*timeout)
74 defer cancel()
75 ch := d.DownloadWithContext(ctx, s.URL)
76 <-ch // discard the first status (file size)
77 got := <-ch
78 if got.Error == nil {
79 t.Error("expected an error due to context timeout, but got nil")
80 }
81 if _, ok := <-ch; ok {
82 t.Error("expected channel closed, but did not get it")
83 }
84}
85
86func TestDownload_DefaultDownloader(t *testing.T) {
87 t.Parallel()
88 s := httptest.NewServer(http.HandlerFunc(
89 func(w http.ResponseWriter, r *http.Request) {
90 if r.Method == http.MethodHead {
91 w.Header().Add("Content-Length", "2")
92 return
93 }
94 if _, err := io.WriteString(w, "42"); err != nil {
95 t.Errorf("failed to write response: %v", err)
96 }
97 },
98 ))
99 defer s.Close()
100
101 d := DefaultDownloader()
102 d.OutputDir = t.TempDir()
103 ch := d.Download(s.URL + "/My%20File.txt")
104 <-ch // discard the first status (just the file size)
105 got := <-ch
106 defer func() {
107 if err := os.Remove(got.DownloadedFilePath); err != nil {
108 t.Errorf("failed to remove test file: %v", err)
109 }
110 }()
111
112 if got.Error != nil {
113 t.Errorf("invalid error. want:nil got:%q", got.Error)
114 }
115 if got.URL != s.URL+"/My%20File.txt" {
116 t.Errorf("invalid URL. want:%s got:%s", s.URL+"/My%20File.txt", got.URL)
117 }
118 if got.DownloadedFileBytes != 2 {
119 t.Errorf("invalid DownloadedFileBytes. want:2 got:%d", got.DownloadedFileBytes)
120 }
121 if !strings.HasSuffix(got.DownloadedFilePath, "My File.txt") {
122 t.Errorf("expected DownloadedFilePath to enfd with My File.txt, got %s", got.DownloadedFilePath)
123 }
124 if got.FileSizeBytes != 2 {
125 t.Errorf("invalid FileSizeBytes. want:2 got:%d", got.FileSizeBytes)
126 }
127 b, err := os.ReadFile(got.DownloadedFilePath)
128 if err != nil {
129 t.Errorf("error reading downloaded file (%s): %q", got.DownloadedFilePath, err)
130 }
131 if string(b) != "42" {
132 t.Errorf("invalid downloaded file content. want:42 got:%s", string(b))
133 }
134 if _, ok := <-ch; ok {
135 t.Error("expected channel closed, but did not get it")
136 }
137}
138
139func TestDownload_ZIPArchive(t *testing.T) {
140 t.Parallel()
141 tmp := t.TempDir()
142 pth := filepath.Join(tmp, "archive.zip")
143 expected := make([]byte, 2<<20)
144 for i := range 2 << 20 {
145 expected[i] = byte(97 + rand.Intn(122-97))
146 }
147
148 // create a zip archive
149 func() {
150 z, err := os.Create(pth)
151 if err != nil {
152 t.Errorf("expected no error creating zip archive, got %s", err)
153 }
154 defer func() {
155 if err := z.Close(); err != nil {
156 t.Errorf("failed to close zip file: %v", err)
157 }
158 }()
159 w := zip.NewWriter(z)
160 f, err := w.Create("file.txt")
161 if err != nil {
162 t.Errorf("expected no error creating archived file, got %s", err)
163 }
164 defer func() {
165 if err := w.Close(); err != nil {
166 t.Errorf("failed to close zip writer: %v", err)
167 }
168 }()
169 if _, err := f.Write(expected); err != nil {
170 t.Errorf("expected no error writing to archived file, got %s", err)
171 }
172 }()
173
174 // create a server to serve the zip archive
175 s := httptest.NewServer(http.HandlerFunc(
176 func(w http.ResponseWriter, r *http.Request) {
177 http.ServeFile(w, r, pth)
178 },
179 ))
180 defer s.Close()
181
182 // download
183 var got string
184 defer func() {
185 if got != "" {
186 if err := os.Remove(got); err != nil {
187 t.Errorf("failed to remove test file: %v", err)
188 }
189 }
190 }()
191 d := DefaultDownloader()
192 d.OutputDir = t.TempDir()
193 for g := range d.Download(s.URL + "/archive.zip") {
194 got = g.DownloadedFilePath
195 if g.Error != nil {
196 t.Errorf("expected no error during the download of the zip archive, got %s", g.Error)
197 }
198 }
199
200 // unarchive and check contents
201 a, err := zip.OpenReader(got)
202 if err != nil {
203 t.Errorf("expected no error opening downloaded zip archive %s, got %s", got, err)
204 }
205 defer func() {
206 if err := a.Close(); err != nil {
207 t.Errorf("failed to close zip reader: %v", err)
208 }
209 }()
210 r, err := a.Open("file.txt")
211 if err != nil {
212 t.Errorf("expected no error reading downloaded zip archive, got %s", err)
213 }
214 defer func() {
215 if err := r.Close(); err != nil {
216 t.Errorf("failed to close zip file reader: %v", err)
217 }
218 }()
219 var b bytes.Buffer
220 if _, err := io.Copy(&b, r); err != nil {
221 t.Errorf("expected no error reading archived file, got %s", err)
222 }
223 if !bytes.Equal(expected, b.Bytes()) {
224 t.Error("archived contents differ from expected") // not printing because it's a lot of data
225 }
226}
227
228func TestDownload_Retry(t *testing.T) {
229 t.Parallel()
230 var done atomic.Bool
231 s := httptest.NewServer(http.HandlerFunc(
232 func(w http.ResponseWriter, r *http.Request) {
233 if r.Method == http.MethodHead {
234 w.Header().Add("Content-Length", "2")
235 return
236 }
237 if done.CompareAndSwap(false, true) {
238 w.WriteHeader(http.StatusTeapot)
239 return
240 }
241 if _, err := io.WriteString(w, "42"); err != nil {
242 t.Errorf("failed to write response: %v", err)
243 }
244 },
245 ))
246 defer s.Close()
247
248 d := Downloader{
249 OutputDir: t.TempDir(),
250 Timeout: timeout,
251 MaxRetries: 4,
252 ConcurrencyPerServer: 1,
253 ChunkSize: 1024,
254 WaitRetry: 1 * time.Millisecond,
255 }
256 ch := d.Download(s.URL)
257 <-ch // discard the first status (just the file size)
258 got := <-ch
259 if got.Error != nil {
260 t.Errorf("invalid error. want:nil got:%q", got.Error)
261 }
262 if !done.Load() {
263 t.Error("expected server to be done, it is not")
264 }
265 if got.URL != s.URL {
266 t.Errorf("invalid URL. want:%s got:%s", s.URL, got.URL)
267 }
268 if got.DownloadedFileBytes != 2 {
269 t.Errorf("invalid DownloadedFileBytes. want:2 got:%d", got.DownloadedFileBytes)
270 }
271 if got.FileSizeBytes != 2 {
272 t.Errorf("invalid FileSizeBytes. want:2 got:%d", got.FileSizeBytes)
273 }
274 b, err := os.ReadFile(got.DownloadedFilePath)
275 if err != nil {
276 t.Errorf("error reading downloaded file (%s): %q", got.DownloadedFilePath, err)
277 }
278 if string(b) != "42" {
279 t.Errorf("invalid downloaded file content. want:42 got:%s", string(b))
280 }
281 if _, ok := <-ch; ok {
282 t.Error("expected channel closed, but did not get it")
283 }
284}
285
286func TestDownload_ReportPreviouslyDownloadedBytes(t *testing.T) {
287 t.Parallel()
288 tmp := t.TempDir()
289 pdir := t.TempDir()
290
291 // server that serves first chunk but always fails on second chunk
292 s := httptest.NewServer(http.HandlerFunc(
293 func(w http.ResponseWriter, r *http.Request) {
294 if r.Method == http.MethodHead {
295 w.Header().Add("Content-Length", "4")
296 return
297 }
298 rangeHeader := r.Header.Get("Range")
299 if strings.HasPrefix(rangeHeader, "bytes=0-") {
300 w.WriteHeader(http.StatusPartialContent)
301 if _, err := io.WriteString(w, "42"); err != nil {
302 t.Errorf("failed to write response: %v", err)
303 }
304 return
305 }
306 w.WriteHeader(http.StatusTeapot)
307 },
308 ))
309 defer s.Close()
310
311 // first download attempt: first chunk succeeds, second chunk always fails
312 d := Downloader{
313 OutputDir: tmp,
314 ProgressDir: pdir,
315 Timeout: timeout,
316 MaxRetries: 1,
317 ConcurrencyPerServer: 1,
318 ChunkSize: 2,
319 }
320 ch := d.Download(s.URL)
321 var first DownloadStatus
322 for s := range ch {
323 first = s
324 }
325
326 // second download attempt: should report the previously downloaded bytes
327 d = Downloader{
328 OutputDir: tmp,
329 ProgressDir: pdir,
330 Timeout: timeout,
331 MaxRetries: 1,
332 ConcurrencyPerServer: 1,
333 ChunkSize: 2,
334 }
335 ch = d.Download(s.URL)
336 var second DownloadStatus
337 for s := range ch {
338 second = s
339 }
340 if first.DownloadedFileBytes != second.DownloadedFileBytes {
341 t.Errorf("expected the same number of downloaded bytes, got %d and %d", first.DownloadedFileBytes, second.DownloadedFileBytes)
342 }
343}
344
345func TestDownloadWithContext_ErrorUserTimeout(t *testing.T) {
346 t.Parallel()
347 usrTimeout := 250 * time.Millisecond // smaller than overall timeout
348 chunkTimeout := 10 * usrTimeout
349 s := httptest.NewServer(http.HandlerFunc(
350 func(w http.ResponseWriter, r *http.Request) {
351 if r.Method == http.MethodHead {
352 w.Header().Add("Content-Length", "2")
353 return
354 }
355 time.Sleep(2 * usrTimeout) // greater than the user timeout, but shorter than the timeout per chunk.
356 },
357 ))
358 defer s.Close()
359 d := Downloader{
360 OutputDir: t.TempDir(),
361 Timeout: chunkTimeout,
362 MaxRetries: 4,
363 ConcurrencyPerServer: 1,
364 ChunkSize: 1024,
365 WaitRetry: 0 * time.Second,
366 }
367 userCtx, cancFunc := context.WithTimeout(context.Background(), usrTimeout)
368 defer cancFunc()
369
370 ch := d.DownloadWithContext(userCtx, s.URL)
371 <-ch // discard the first got (just the file size)
372 got := <-ch
373 if got.Error == nil {
374 t.Error("expected an error, but got nil")
375 }
376 if _, ok := <-ch; ok {
377 t.Error("expected channel closed, but did not get it")
378 }
379}
380
381func TestDownload_Chunks(t *testing.T) {
382 t.Parallel()
383 d := DefaultDownloader()
384 d.ChunkSize = 5
385 got := d.chunks(12)
386 chunks := []chunk{{0, 4}, {5, 9}, {10, 11}}
387 sizes := []int64{5, 5, 2}
388 headers := []string{"bytes=0-4", "bytes=5-9", "bytes=10-11"}
389 if len(got) != len(chunks) {
390 t.Errorf("expected %d chunks, got %d", len(chunks), len(got))
391 }
392 for i := range got {
393 if got[i].start != chunks[i].start {
394 t.Errorf("expected chunk #%d to start at %d, got %d", i+1, chunks[i].start, got[i].start)
395 }
396 if got[i].end != chunks[i].end {
397 t.Errorf("expected chunk #%d to end at %d, got %d", i+1, chunks[i].end, got[i].end)
398 }
399 if got[i].size() != sizes[i] {
400 t.Errorf("expected chunk #%d to have size %d, got %d", i+1, sizes[i], got[i].size())
401 }
402 if got[i].rangeHeader() != headers[i] {
403 t.Errorf("expected chunk #%d header to be %s, got %s", i+1, headers[i], got[i].rangeHeader())
404 }
405 }
406}
407
408func TestGetDownload_WithUserAgent(t *testing.T) {
409 t.Parallel()
410 ua := "Answer/42.0"
411 s := httptest.NewServer(http.HandlerFunc(
412 func(w http.ResponseWriter, r *http.Request) {
413 if r.UserAgent() != ua {
414 t.Errorf("expected user-agent to be %s, got %s", ua, r.UserAgent())
415 }
416 },
417 ))
418 defer s.Close()
419 d := DefaultDownloader()
420 d.UserAgent = ua
421 <-d.Download(s.URL)
422}
423
424func TestGetDownloadSize_ContentLength(t *testing.T) {
425 t.Parallel()
426 s := httptest.NewServer(http.HandlerFunc(
427 func(w http.ResponseWriter, r *http.Request) {
428 if _, err := io.WriteString(w, "Test"); err != nil {
429 t.Errorf("failed to write response: %v", err)
430 }
431 },
432 ))
433 defer s.Close()
434
435 d := DefaultDownloader()
436 got, err := d.getDownloadSize(context.Background(), s.URL)
437
438 if err != nil {
439 t.Errorf("expected no error getting the file size, got %s", err)
440 }
441 if got != 4 {
442 t.Errorf("invalid size, expected 4, got: %d", got)
443 }
444}
445
446func TestGetDownloadSize_WithRetry(t *testing.T) {
447 t.Parallel()
448 attempts := int32(0)
449 s := httptest.NewServer(http.HandlerFunc(
450 func(w http.ResponseWriter, r *http.Request) {
451 if atomic.CompareAndSwapInt32(&attempts, 0, 1) {
452 w.WriteHeader(http.StatusTooManyRequests)
453 return
454 }
455 if _, err := io.WriteString(w, "Test"); err != nil {
456 t.Errorf("failed to write response: %v", err)
457 }
458 },
459 ))
460 defer s.Close()
461
462 d := DefaultDownloader()
463 got, err := d.getDownloadSize(context.Background(), s.URL)
464
465 if err != nil {
466 t.Errorf("expected no error getting the file size, got %s", err)
467 }
468 if got != 4 {
469 t.Errorf("invalid size, expected 4, got: %d", got)
470 }
471}
472
473func TestGetDownloadSize_ContentRange(t *testing.T) {
474 t.Parallel()
475 s := httptest.NewServer(http.HandlerFunc(
476 func(w http.ResponseWriter, r *http.Request) {
477 w.Header().Set("Content-Range", "bytes 1-10/123")
478 if _, err := io.WriteString(w, ""); err != nil {
479 t.Errorf("failed to write response: %v", err)
480 }
481 },
482 ))
483 defer s.Close()
484
485 d := DefaultDownloader()
486 got, err := d.getDownloadSize(context.Background(), s.URL)
487
488 if err != nil {
489 t.Errorf("expected no error getting the file size, got %s", err)
490 }
491 if got != 123 {
492 t.Errorf("invalid size, expected 123, got: %d", got)
493 }
494}
495
496func TestGetDownloadSize_ErrorInvalidURL(t *testing.T) {
497 t.Parallel()
498 d := Downloader{
499 MaxRetries: 1,
500 WaitRetry: 0,
501 Client: http.DefaultClient,
502 }
503 got, err := d.getDownloadSize(context.Background(), "test")
504
505 if err == nil {
506 t.Errorf("expected an error, got nil")
507 }
508 if got != 0 {
509 t.Errorf("invalid size, expected 0, got: %d", got)
510 }
511}
512
513func TestGetDownloadSize_NoContent(t *testing.T) {
514 t.Parallel()
515 s := httptest.NewServer(http.HandlerFunc(
516 func(w http.ResponseWriter, r *http.Request) {
517 if _, err := io.WriteString(w, ""); err != nil {
518 t.Errorf("failed to write response: %v", err)
519 }
520 },
521 ))
522 defer s.Close()
523
524 d := DefaultDownloader()
525 if _, err := d.getDownloadSize(context.Background(), s.URL); err == nil {
526 t.Error("expected error getting the file size, got nil")
527 }
528}