Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package upload
18

19
import (
20
	"bytes"
21
	"context"
22
	"fmt"
23
	"io"
24
	"net"
25
	"net/http"
26
	"os"
27
	"testing"
28

29
	testifyassert "github.com/stretchr/testify/assert"
30
	"go.uber.org/mock/gomock"
31
	"golang.org/x/time/rate"
32

33
	"d7y.io/dragonfly/v2/client/config"
34
	"d7y.io/dragonfly/v2/client/daemon/storage"
35
	"d7y.io/dragonfly/v2/client/daemon/storage/mocks"
36
	"d7y.io/dragonfly/v2/client/daemon/test"
37
	_ "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/server"
38
)
39

40
func TestUploadManager_Serve(t *testing.T) {
41
	ctrl := gomock.NewController(t)
42
	defer ctrl.Finish()
43

44
	assert := testifyassert.New(t)
45
	testData, err := os.ReadFile(test.File)
46
	assert.Nil(err, "load test file")
47

48
	mockStorageManager := mocks.NewMockManager(ctrl)
49
	mockStorageManager.EXPECT().ReadPiece(gomock.Any(), gomock.Any()).AnyTimes().
50
		DoAndReturn(func(ctx context.Context, req *storage.ReadPieceRequest) (io.Reader, io.Closer, error) {
51
			return bytes.NewBuffer(testData[req.Range.Start : req.Range.Start+req.Range.Length]),
52
				io.NopCloser(nil), nil
53
		})
54

55
	um, err := NewUploadManager(config.NewDaemonConfig(), mockStorageManager, os.TempDir(), WithLimiter(rate.NewLimiter(16*1024, 16*1024)))
56
	assert.Nil(err, "NewUploadManager")
57

58
	listen, err := net.Listen("tcp4", "127.0.0.1:0")
59
	assert.Nil(err, "Listen")
60
	addr := listen.Addr().String()
61

62
	go func() {
63
		if err := um.Serve(listen); err != nil {
64
			t.Error(err)
65
		}
66
	}()
67

68
	tests := []struct {
69
		taskID          string
70
		peerID          string
71
		pieceRange      string
72
		targetPieceData []byte
73
	}{
74
		{
75
			taskID:          "task-0",
76
			peerID:          "peer-0",
77
			pieceRange:      "bytes=0-9",
78
			targetPieceData: testData[0:10],
79
		},
80
		{
81
			taskID:          "task-1",
82
			peerID:          "peer-1",
83
			pieceRange:      fmt.Sprintf("bytes=512-%d", len(testData)-1),
84
			targetPieceData: testData[512:],
85
		},
86
		{
87
			taskID:          "task-2",
88
			peerID:          "peer-2",
89
			pieceRange:      "bytes=512-1023",
90
			targetPieceData: testData[512:1024],
91
		},
92
	}
93

94
	for _, tt := range tests {
95
		req, _ := http.NewRequest(http.MethodGet,
96
			fmt.Sprintf("http://%s/%s/%s/%s?peerId=%s", addr, "download", "666", tt.taskID, tt.peerID), nil)
97
		req.Header.Add("Range", tt.pieceRange)
98

99
		resp, err := http.DefaultClient.Do(req)
100
		assert.Nil(err, "get piece data")
101

102
		data, _ := io.ReadAll(resp.Body)
103
		resp.Body.Close()
104
		assert.Equal(tt.targetPieceData, data)
105
	}
106
}
107