github.com/Heebron/moby@v0.0.0-20221111184709-6eab4f55faf7/distribution/xfer/upload_test.go (about)

     1  package xfer // import "github.com/docker/docker/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/docker/docker/layer"
    12  	"github.com/docker/docker/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: currentUploads, diffID: "sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"},
    73  		&mockUploadDescriptor{currentUploads: currentUploads, diffID: "sha256:1515325234325236634634608943609283523908626098235490238423902343"},
    74  		&mockUploadDescriptor{currentUploads: currentUploads, diffID: "sha256:6929356290463485374960346430698374523437683470934634534953453453"},
    75  		&mockUploadDescriptor{currentUploads: currentUploads, diffID: "sha256:cbbf2f9a99b47fc460d422812b6a5adff7dfee951d8fa2e4a98caa0382cfbdbf"},
    76  		&mockUploadDescriptor{currentUploads: currentUploads, diffID: "sha256:8159352387436803946235346346368745389534789534897538734598734987", simulateRetries: 1},
    77  		&mockUploadDescriptor{currentUploads: currentUploads, diffID: "sha256:4637863963478346897346987346987346789346789364879364897364987346"},
    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(&currentUploads)
    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  }