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 transport
18

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

27
	testifyassert "github.com/stretchr/testify/assert"
28
	"go.uber.org/mock/gomock"
29

30
	"d7y.io/dragonfly/v2/client/daemon/peer"
31
	"d7y.io/dragonfly/v2/client/daemon/test"
32
)
33

34
func TestTransport_RoundTrip(t *testing.T) {
35
	assert := testifyassert.New(t)
36
	ctrl := gomock.NewController(t)
37
	testData, err := os.ReadFile(test.File)
38
	assert.Nil(err, "load test file")
39

40
	var url = "http://x/y"
41
	peerTaskManager := peer.NewMockTaskManager(ctrl)
42
	peerTaskManager.EXPECT().StartStreamTask(gomock.Any(), gomock.Any()).DoAndReturn(
43
		func(ctx context.Context, req *peer.StreamTaskRequest) (io.ReadCloser, map[string]string, error) {
44
			assert.Equal(req.URL, url)
45
			return io.NopCloser(bytes.NewBuffer(testData)), nil, nil
46
		},
47
	)
48
	rt, _ := New(
49
		WithPeerIDGenerator(peer.NewPeerIDGenerator("127.0.0.1")),
50
		WithPeerTaskManager(peerTaskManager),
51
		WithCondition(func(r *http.Request) bool {
52
			return true
53
		}))
54
	assert.NotNil(rt)
55
	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
56
	resp, err := rt.RoundTrip(req)
57
	assert.Nil(err)
58
	if err != nil {
59
		return
60
	}
61
	defer resp.Body.Close()
62
	output, err := io.ReadAll(resp.Body)
63
	assert.Nil(err)
64
	if err != nil {
65
		return
66
	}
67
	assert.Equal(testData, output)
68
}
69