github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/distribution/xfer/upload_test.go (about) 1 package xfer // import "github.com/demonoid81/moby/distribution/xfer" 2 3 import ( 4 "context" 5 "errors" 6 "sync/atomic" 7 "testing" 8 "time" 9 10 "github.com/docker/distribution" 11 "github.com/demonoid81/moby/layer" 12 "github.com/demonoid81/moby/pkg/progress" 13 ) 14 15 const maxUploadConcurrency = 3 16 17 type mockUploadDescriptor struct { 18 currentUploads *int32 19 diffID layer.DiffID 20 simulateRetries int 21 } 22 23 // Key returns the key used to deduplicate downloads. 24 func (u *mockUploadDescriptor) Key() string { 25 return u.diffID.String() 26 } 27 28 // ID returns the ID for display purposes. 29 func (u *mockUploadDescriptor) ID() string { 30 return u.diffID.String() 31 } 32 33 // DiffID should return the DiffID for this layer. 34 func (u *mockUploadDescriptor) DiffID() layer.DiffID { 35 return u.diffID 36 } 37 38 // SetRemoteDescriptor is not used in the mock. 39 func (u *mockUploadDescriptor) SetRemoteDescriptor(remoteDescriptor distribution.Descriptor) { 40 } 41 42 // Upload is called to perform the upload. 43 func (u *mockUploadDescriptor) Upload(ctx context.Context, progressOutput progress.Output) (distribution.Descriptor, error) { 44 if u.currentUploads != nil { 45 defer atomic.AddInt32(u.currentUploads, -1) 46 47 if atomic.AddInt32(u.currentUploads, 1) > maxUploadConcurrency { 48 return distribution.Descriptor{}, errors.New("concurrency limit exceeded") 49 } 50 } 51 52 // Sleep a bit to simulate a time-consuming upload. 53 for i := int64(0); i <= 10; i++ { 54 select { 55 case <-ctx.Done(): 56 return distribution.Descriptor{}, ctx.Err() 57 case <-time.After(10 * time.Millisecond): 58 progressOutput.WriteProgress(progress.Progress{ID: u.ID(), Current: i, Total: 10}) 59 } 60 } 61 62 if u.simulateRetries != 0 { 63 u.simulateRetries-- 64 return distribution.Descriptor{}, errors.New("simulating retry") 65 } 66 67 return distribution.Descriptor{}, nil 68 } 69 70 func uploadDescriptors(currentUploads *int32) []UploadDescriptor { 71 return []UploadDescriptor{ 72 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"), 0}, 73 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:1515325234325236634634608943609283523908626098235490238423902343"), 0}, 74 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:6929356290463485374960346430698374523437683470934634534953453453"), 0}, 75 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"), 0}, 76 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:8159352387436803946235346346368745389534789534897538734598734987"), 1}, 77 &mockUploadDescriptor{currentUploads, layer.DiffID("sha256:4637863963478346897346987346987346789346789364879364897364987346"), 0}, 78 } 79 } 80 81 func TestSuccessfulUpload(t *testing.T) { 82 lum := NewLayerUploadManager(maxUploadConcurrency, func(m *LayerUploadManager) { m.waitDuration = time.Millisecond }) 83 84 progressChan := make(chan progress.Progress) 85 progressDone := make(chan struct{}) 86 receivedProgress := make(map[string]int64) 87 88 go func() { 89 for p := range progressChan { 90 receivedProgress[p.ID] = p.Current 91 } 92 close(progressDone) 93 }() 94 95 var currentUploads int32 96 descriptors := uploadDescriptors(¤tUploads) 97 98 err := lum.Upload(context.Background(), descriptors, progress.ChanOutput(progressChan)) 99 if err != nil { 100 t.Fatalf("upload error: %v", err) 101 } 102 103 close(progressChan) 104 <-progressDone 105 } 106 107 func TestCancelledUpload(t *testing.T) { 108 lum := NewLayerUploadManager(maxUploadConcurrency, func(m *LayerUploadManager) { m.waitDuration = time.Millisecond }) 109 110 progressChan := make(chan progress.Progress) 111 progressDone := make(chan struct{}) 112 113 go func() { 114 for range progressChan { 115 } 116 close(progressDone) 117 }() 118 119 ctx, cancel := context.WithCancel(context.Background()) 120 121 go func() { 122 <-time.After(time.Millisecond) 123 cancel() 124 }() 125 126 descriptors := uploadDescriptors(nil) 127 err := lum.Upload(ctx, descriptors, progress.ChanOutput(progressChan)) 128 if err != context.Canceled { 129 t.Fatal("expected upload to be cancelled") 130 } 131 132 close(progressChan) 133 <-progressDone 134 }