test: scheduler supervisor (#742)

Signed-off-by: cutecutecat <12032261@mail.sustech.edu.cn>

scheduler/supervisor/cdn.go

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+//go:generate mockgen -destination ./mocks/cdn_mock.go -package mocks d7y.io/dragonfly/v2/scheduler/supervisor CDNDynmaicClient
 
 package supervisor
 

scheduler/supervisor/cdn_test.go

@@ -0,0 +1,522 @@
+/*
+ *     Copyright 2020 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package supervisor_test
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"reflect"
+	"testing"
+
+	"d7y.io/dragonfly/v2/internal/dfcodes"
+	"d7y.io/dragonfly/v2/internal/dferrors"
+	logger "d7y.io/dragonfly/v2/internal/dflog"
+	"d7y.io/dragonfly/v2/pkg/rpc/base"
+	"d7y.io/dragonfly/v2/pkg/rpc/cdnsystem"
+	"d7y.io/dragonfly/v2/pkg/rpc/cdnsystem/client"
+	"d7y.io/dragonfly/v2/scheduler/supervisor"
+	"d7y.io/dragonfly/v2/scheduler/supervisor/mocks"
+	"github.com/agiledragon/gomonkey"
+	"github.com/golang/mock/gomock"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	mockPieceSeedStream = &client.PieceSeedStream{}
+	mockPieceSeed       = &cdnsystem.PieceSeed{}
+	mockHost            = &supervisor.Host{}
+	mockTask            = &supervisor.Task{}
+	mockPeer            = &supervisor.Peer{}
+	mockLogger          = &logger.SugaredLoggerOnWith{}
+)
+
+func TestCDN_Nil(t *testing.T) {
+	tests := []struct {
+		name   string
+		status supervisor.TaskStatus
+		mock   func(t *testing.T) (supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches)
+		expect func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error)
+	}{
+		{
+			name:   "nil client",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockTask.ID = "mocktask"
+
+				patch := &gomonkey.Patches{}
+				return mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.NotNil(cdn)
+				assert.Nil(cdn.GetClient())
+				assert.Nil(peer)
+				assert.Equal(supervisor.ErrCDNClientUninitialized, err)
+
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			mockPeerManager, mockHostManager, patch := tc.mock(t)
+			cdn := supervisor.NewCDN(nil, mockPeerManager, mockHostManager)
+			mockTask.SetStatus(tc.status)
+			peer, err := cdn.StartSeedTask(context.Background(), mockTask)
+			tc.expect(t, cdn, peer, err)
+			patch.Reset()
+		})
+	}
+}
+
+func TestCDN_Initial(t *testing.T) {
+	tests := []struct {
+		name   string
+		task   *supervisor.Task
+		status supervisor.TaskStatus
+		mock   func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches)
+		expect func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error)
+	}{
+		{
+			name:   "ObtainSeeds cause CdnTaskRegistryFail",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				err := dferrors.New(dfcodes.CdnTaskRegistryFail, "mockError")
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(nil, err).AnyTimes()
+
+				patch := &gomonkey.Patches{}
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNRegisterFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "ObtainSeeds cause CdnTaskDownloadFail",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				err := dferrors.New(dfcodes.CdnTaskDownloadFail, "mockError")
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(nil, err).AnyTimes()
+
+				patch := &gomonkey.Patches{}
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNDownloadFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "ObtainSeeds cause other errors",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				err := dferrors.New(114514, "mockError")
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(nil, err).AnyTimes()
+
+				patch := &gomonkey.Patches{}
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNUnknown, errors.Cause(err))
+			},
+		},
+		{
+			name:   "ObtainSeeds cause invoke client failed",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				err := fmt.Errorf("invoke error")
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(nil, err).AnyTimes()
+
+				patch := &gomonkey.Patches{}
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Equal(supervisor.ErrCDNInvokeFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "failed for EOF and TaskStatusWaiting",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, io.EOF}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(err)
+			},
+		},
+		{
+			name:   "success for EOF and TaskStatusSuccess",
+			status: supervisor.TaskStatusSuccess,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+				mockPeerManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+				mockPeerManager.EXPECT().Add(gomock.Any()).Return().AnyTimes()
+				mockHostManager.EXPECT().Get(gomock.Any()).Return(mockHost, true).AnyTimes()
+
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockLogger), "Debugf",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockTask), "GetOrAddPiece",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{nil, true}}})
+				patch.ApplyMethodSeq(reflect.TypeOf(mockTask), "Log",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{mockLogger}}})
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPeer), "Touch",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPeer), "UpdateProgress",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				newPeerRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPeer}},
+				}
+				patch.ApplyFuncSeq(supervisor.NewPeer, newPeerRet)
+
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPieceSeed, nil}, Times: 1},
+					{Values: gomonkey.Params{nil, io.EOF}, Times: 1},
+				}
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+
+				mockPieceSeed.PieceInfo = &base.PieceInfo{PieceNum: 0}
+				mockPeer.Task = mockTask
+				mockPeer.ID = "114514"
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Equal(mockPeer, peer)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "receivePiece cause CdnTaskRegistryFail",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+
+				err := dferrors.New(dfcodes.CdnTaskRegistryFail, "mockError")
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, err}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNRegisterFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "receivePiece cause CdnTaskDownloadFail",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+
+				err := dferrors.New(dfcodes.CdnTaskDownloadFail, "mockError")
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, err}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNDownloadFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "receivePiece cause other errors",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+
+				err := dferrors.New(114514, "mockError")
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, err}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Error(supervisor.ErrCDNUnknown, errors.Cause(err))
+			},
+		},
+		{
+			name: "receivePiece cause invoke client failed",
+
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+
+				err := fmt.Errorf("invoke error")
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, err}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Equal(supervisor.ErrCDNInvokeFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "initCDNPeer peer is nil",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+				mockCDNDynmaicClient.EXPECT().GetHost(gomock.Any()).Return(nil, false).AnyTimes()
+				mockPeerManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+				mockHostManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPieceSeed, nil}},
+				}
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.Equal(supervisor.ErrInitCDNPeerFail, errors.Cause(err))
+			},
+		},
+		{
+			name:   "downloadTinyFile http.Get error (restore host from hostManager)",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+				mockCDNDynmaicClient.EXPECT().GetHost(gomock.Any()).Return(nil, false).AnyTimes()
+				mockPeerManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+				mockPeerManager.EXPECT().Add(gomock.Any()).Return().AnyTimes()
+				mockHostManager.EXPECT().Get(gomock.Any()).Return(mockHost, true).AnyTimes()
+
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockLogger), "Debugf",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockTask), "Log",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{mockLogger}}})
+
+				err := fmt.Errorf("http error")
+				httpRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{nil, err}},
+				}
+
+				patch.ApplyFuncSeq(http.Get, httpRet)
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPeer), "Touch",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				newPeerRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPeer}},
+				}
+				patch.ApplyFuncSeq(supervisor.NewPeer, newPeerRet)
+
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPieceSeed, nil}},
+				}
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+
+				mockPieceSeed.Done = true
+				mockHost.IP = "0.0.0.0"
+				mockHost.DownloadPort = 1919
+				mockTask.ID = "1919810"
+				mockPeer.Host = mockHost
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Equal(mockPeer, peer)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "downloadTinyFile success (restore host from client)",
+			status: supervisor.TaskStatusWaiting,
+			mock: func(t *testing.T) (supervisor.CDNDynmaicClient, supervisor.PeerManager, supervisor.HostManager, *gomonkey.Patches) {
+				ctl := gomock.NewController(t)
+				defer ctl.Finish()
+
+				mockCDNDynmaicClient := mocks.NewMockCDNDynmaicClient(ctl)
+				mockPeerManager := mocks.NewMockPeerManager(ctl)
+				mockHostManager := mocks.NewMockHostManager(ctl)
+				mockCDNDynmaicClient.EXPECT().ObtainSeeds(gomock.Any(), gomock.Any()).Return(mockPieceSeedStream, nil).AnyTimes()
+				mockCDNDynmaicClient.EXPECT().GetHost(gomock.Any()).Return(mockHost, true).AnyTimes()
+				mockPeerManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+				mockPeerManager.EXPECT().Add(gomock.Any()).Return().AnyTimes()
+				mockHostManager.EXPECT().Get(gomock.Any()).Return(nil, false).AnyTimes()
+				mockHostManager.EXPECT().Add(gomock.Any()).Return().AnyTimes()
+
+				patch := gomonkey.ApplyMethodSeq(reflect.TypeOf(mockLogger), "Debugf",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockTask), "Log",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{mockLogger}}})
+
+				const testwords string = "dragonfly-scheduler-test"
+				res := &http.Response{
+					Body: ioutil.NopCloser(
+						bytes.NewBuffer([]byte(testwords))),
+				}
+				httpRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{res, nil}},
+				}
+				patch.ApplyFuncSeq(http.Get, httpRet)
+
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPeer), "Touch",
+					[]gomonkey.OutputCell{{Values: gomonkey.Params{}}})
+
+				newPeerRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPeer}},
+				}
+				patch.ApplyFuncSeq(supervisor.NewPeer, newPeerRet)
+
+				streamRet := []gomonkey.OutputCell{
+					{Values: gomonkey.Params{mockPieceSeed, nil}},
+				}
+				patch.ApplyMethodSeq(reflect.TypeOf(mockPieceSeedStream), "Recv", streamRet)
+
+				mockPieceSeed.Done = true
+				mockPieceSeed.ContentLength = int64(len(testwords))
+				mockHost.IP = "0.0.0.0"
+				mockHost.DownloadPort = 1919
+				mockTask.ID = "1919810"
+				mockPeer.Host = mockHost
+				return mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch
+			},
+			expect: func(t *testing.T, cdn supervisor.CDN, peer *supervisor.Peer, err error) {
+				assert := assert.New(t)
+				assert.Equal(mockPeer, peer)
+				assert.Nil(err)
+				assert.Equal([]byte("dragonfly-scheduler-test"), mockTask.DirectPiece)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			mockCDNDynmaicClient, mockPeerManager, mockHostManager, patch := tc.mock(t)
+			cdn := supervisor.NewCDN(mockCDNDynmaicClient, mockPeerManager, mockHostManager)
+			mockTask.SetStatus(tc.status)
+			peer, err := cdn.StartSeedTask(context.Background(), mockTask)
+			tc.expect(t, cdn, peer, err)
+			patch.Reset()
+		})
+	}
+}

scheduler/supervisor/host.go

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+//go:generate mockgen -destination ./mocks/host_mock.go -package mocks d7y.io/dragonfly/v2/scheduler/supervisor HostManager
+
 package supervisor
 
 import (

scheduler/supervisor/host_test.go

@@ -0,0 +1,221 @@
+/*
+ *     Copyright 2020 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package supervisor_test
+
+import (
+	"strconv"
+	"testing"
+
+	"d7y.io/dragonfly/v2/scheduler/supervisor"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHost_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		host   *supervisor.Host
+		expect func(t *testing.T, host *supervisor.Host)
+	}{
+		{
+			name: "create by normal config",
+			host: supervisor.NewClientHost("main", "127.0.0.1", "Client", 8080, 8081, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("main", host.UUID)
+			},
+		},
+		{
+			name: "create CDN by normal config",
+			host: supervisor.NewCDNHost("main", "127.0.0.1", "Client", 8080, 8081, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("main", host.UUID)
+			},
+		},
+		{
+			name: "create by special symbols",
+			host: supervisor.NewClientHost("⁂⁎♜♝♞⁑（๑ `▽´๑)", "127.0.0.1", "Client", 8080, 8081, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("⁂⁎♜♝♞⁑（๑ `▽´๑)", host.UUID)
+			},
+		},
+		{
+			name: "create by error address",
+			host: supervisor.NewClientHost("host", "0.0.0.0", "Client", 8080, 8080, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("host", host.UUID)
+			},
+		},
+		{
+			name: "create with geography information",
+			host: supervisor.NewClientHost("host", "127.0.0.1", "Client", 8080, 8081, "goagle", "microsaft", "facebaok"),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("host", host.UUID)
+				assert.Equal("goagle", host.SecurityDomain)
+				assert.Equal("microsaft", host.Location)
+				assert.Equal("facebaok", host.IDC)
+			},
+		},
+		{
+			name: "create by error address",
+			host: supervisor.NewClientHost("host", "-1.257.w.-0", "Client", -100, 29000, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("host", host.UUID)
+			},
+		},
+		{
+			name: "create by normal config",
+			host: supervisor.NewClientHost("host", "127.0.0.1", "Client", 8080, 8081, "", "", ""),
+			expect: func(t *testing.T, host *supervisor.Host) {
+				assert := assert.New(t)
+				assert.Equal("host", host.UUID)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.expect(t, tc.host)
+		})
+	}
+}
+
+func TestHostManager_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		expect func(t *testing.T, hostManager supervisor.HostManager)
+	}{
+		{
+			name: "simple create",
+			expect: func(t *testing.T, hostManager supervisor.HostManager) {
+				assert := assert.New(t)
+				assert.NotNil(hostManager)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			hostManager := supervisor.NewHostManager()
+			tc.expect(t, hostManager)
+		})
+	}
+}
+
+func TestHostManager_Get(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		fetch  int
+		expect func(t *testing.T, host *supervisor.Host, success bool)
+	}{
+		{
+			name:   "fetch first host",
+			number: 3,
+			fetch:  0,
+			expect: func(t *testing.T, host *supervisor.Host, success bool) {
+				assert := assert.New(t)
+				assert.Equal("0", host.UUID)
+				assert.True(success)
+			},
+		},
+		{
+			name:   "fetch last host",
+			number: 3,
+			fetch:  2,
+			expect: func(t *testing.T, host *supervisor.Host, success bool) {
+				assert := assert.New(t)
+				assert.Equal("2", host.UUID)
+				assert.True(success)
+			},
+		},
+		{
+			name:   "fetch not exist host",
+			number: 3,
+			fetch:  -1,
+			expect: func(t *testing.T, host *supervisor.Host, success bool) {
+				assert := assert.New(t)
+				assert.Equal((*supervisor.Host)(nil), host)
+				assert.False(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			hostManager := supervisor.NewHostManager()
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				host := mockAHost(index)
+				hostManager.Add(host)
+			}
+			host, success := hostManager.Get(strconv.Itoa(tc.fetch))
+			tc.expect(t, host, success)
+		})
+	}
+}
+
+func TestHostManager_Delete(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		delete int
+		expect func(t *testing.T, host *supervisor.Host, success bool)
+	}{
+		{
+			name:   "delete exist host",
+			number: 1,
+			delete: 0,
+			expect: func(t *testing.T, host *supervisor.Host, success bool) {
+				assert := assert.New(t)
+				assert.Nil(host)
+				assert.False(success)
+			},
+		},
+		{
+			name:   "delete not exist host",
+			number: 1,
+			delete: 100,
+			expect: func(t *testing.T, host *supervisor.Host, success bool) {
+				assert := assert.New(t)
+				assert.Nil(host)
+				assert.False(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			hostManager := supervisor.NewHostManager()
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				host := mockAHost(index)
+				hostManager.Add(host)
+			}
+			hostManager.Delete(strconv.Itoa(tc.delete))
+			host, success := hostManager.Get(strconv.Itoa(tc.delete))
+
+			tc.expect(t, host, success)
+		})
+	}
+}
+
+func mockAHost(UUID string) *supervisor.Host {
+	host := supervisor.NewClientHost(UUID, "127.0.0.1", "Client", 8080, 8081, "", "", "")
+	return host
+}

scheduler/supervisor/mocks/cdn_mock.go

@@ -0,0 +1,135 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: d7y.io/dragonfly/v2/scheduler/supervisor (interfaces: CDNDynmaicClient)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	context "context"
+	reflect "reflect"
+
+	dfnet "d7y.io/dragonfly/v2/pkg/basic/dfnet"
+	base "d7y.io/dragonfly/v2/pkg/rpc/base"
+	cdnsystem "d7y.io/dragonfly/v2/pkg/rpc/cdnsystem"
+	client "d7y.io/dragonfly/v2/pkg/rpc/cdnsystem/client"
+	config "d7y.io/dragonfly/v2/scheduler/config"
+	supervisor "d7y.io/dragonfly/v2/scheduler/supervisor"
+	gomock "github.com/golang/mock/gomock"
+	grpc "google.golang.org/grpc"
+)
+
+// MockCDNDynmaicClient is a mock of CDNDynmaicClient interface.
+type MockCDNDynmaicClient struct {
+	ctrl     *gomock.Controller
+	recorder *MockCDNDynmaicClientMockRecorder
+}
+
+// MockCDNDynmaicClientMockRecorder is the mock recorder for MockCDNDynmaicClient.
+type MockCDNDynmaicClientMockRecorder struct {
+	mock *MockCDNDynmaicClient
+}
+
+// NewMockCDNDynmaicClient creates a new mock instance.
+func NewMockCDNDynmaicClient(ctrl *gomock.Controller) *MockCDNDynmaicClient {
+	mock := &MockCDNDynmaicClient{ctrl: ctrl}
+	mock.recorder = &MockCDNDynmaicClientMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockCDNDynmaicClient) EXPECT() *MockCDNDynmaicClientMockRecorder {
+	return m.recorder
+}
+
+// Close mocks base method.
+func (m *MockCDNDynmaicClient) Close() error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Close")
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Close indicates an expected call of Close.
+func (mr *MockCDNDynmaicClientMockRecorder) Close() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockCDNDynmaicClient)(nil).Close))
+}
+
+// GetHost mocks base method.
+func (m *MockCDNDynmaicClient) GetHost(arg0 string) (*supervisor.Host, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetHost", arg0)
+	ret0, _ := ret[0].(*supervisor.Host)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// GetHost indicates an expected call of GetHost.
+func (mr *MockCDNDynmaicClientMockRecorder) GetHost(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetHost", reflect.TypeOf((*MockCDNDynmaicClient)(nil).GetHost), arg0)
+}
+
+// GetPieceTasks mocks base method.
+func (m *MockCDNDynmaicClient) GetPieceTasks(arg0 context.Context, arg1 dfnet.NetAddr, arg2 *base.PieceTaskRequest, arg3 ...grpc.CallOption) (*base.PiecePacket, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1, arg2}
+	for _, a := range arg3 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetPieceTasks", varargs...)
+	ret0, _ := ret[0].(*base.PiecePacket)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetPieceTasks indicates an expected call of GetPieceTasks.
+func (mr *MockCDNDynmaicClientMockRecorder) GetPieceTasks(arg0, arg1, arg2 interface{}, arg3 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1, arg2}, arg3...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPieceTasks", reflect.TypeOf((*MockCDNDynmaicClient)(nil).GetPieceTasks), varargs...)
+}
+
+// ObtainSeeds mocks base method.
+func (m *MockCDNDynmaicClient) ObtainSeeds(arg0 context.Context, arg1 *cdnsystem.SeedRequest, arg2 ...grpc.CallOption) (*client.PieceSeedStream, error) {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "ObtainSeeds", varargs...)
+	ret0, _ := ret[0].(*client.PieceSeedStream)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ObtainSeeds indicates an expected call of ObtainSeeds.
+func (mr *MockCDNDynmaicClientMockRecorder) ObtainSeeds(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ObtainSeeds", reflect.TypeOf((*MockCDNDynmaicClient)(nil).ObtainSeeds), varargs...)
+}
+
+// OnNotify mocks base method.
+func (m *MockCDNDynmaicClient) OnNotify(arg0 *config.DynconfigData) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "OnNotify", arg0)
+}
+
+// OnNotify indicates an expected call of OnNotify.
+func (mr *MockCDNDynmaicClientMockRecorder) OnNotify(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnNotify", reflect.TypeOf((*MockCDNDynmaicClient)(nil).OnNotify), arg0)
+}
+
+// UpdateState mocks base method.
+func (m *MockCDNDynmaicClient) UpdateState(arg0 []dfnet.NetAddr) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "UpdateState", arg0)
+}
+
+// UpdateState indicates an expected call of UpdateState.
+func (mr *MockCDNDynmaicClientMockRecorder) UpdateState(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateState", reflect.TypeOf((*MockCDNDynmaicClient)(nil).UpdateState), arg0)
+}

scheduler/supervisor/mocks/gc_mock.go

@@ -0,0 +1,99 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: d7y.io/dragonfly/v2/pkg/gc (interfaces: GC)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	reflect "reflect"
+
+	gc "d7y.io/dragonfly/v2/pkg/gc"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockGC is a mock of GC interface.
+type MockGC struct {
+	ctrl     *gomock.Controller
+	recorder *MockGCMockRecorder
+}
+
+// MockGCMockRecorder is the mock recorder for MockGC.
+type MockGCMockRecorder struct {
+	mock *MockGC
+}
+
+// NewMockGC creates a new mock instance.
+func NewMockGC(ctrl *gomock.Controller) *MockGC {
+	mock := &MockGC{ctrl: ctrl}
+	mock.recorder = &MockGCMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockGC) EXPECT() *MockGCMockRecorder {
+	return m.recorder
+}
+
+// Add mocks base method.
+func (m *MockGC) Add(arg0 gc.Task) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Add", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Add indicates an expected call of Add.
+func (mr *MockGCMockRecorder) Add(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Add", reflect.TypeOf((*MockGC)(nil).Add), arg0)
+}
+
+// Run mocks base method.
+func (m *MockGC) Run(arg0 string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Run", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Run indicates an expected call of Run.
+func (mr *MockGCMockRecorder) Run(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Run", reflect.TypeOf((*MockGC)(nil).Run), arg0)
+}
+
+// RunAll mocks base method.
+func (m *MockGC) RunAll() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "RunAll")
+}
+
+// RunAll indicates an expected call of RunAll.
+func (mr *MockGCMockRecorder) RunAll() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RunAll", reflect.TypeOf((*MockGC)(nil).RunAll))
+}
+
+// Serve mocks base method.
+func (m *MockGC) Serve() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Serve")
+}
+
+// Serve indicates an expected call of Serve.
+func (mr *MockGCMockRecorder) Serve() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Serve", reflect.TypeOf((*MockGC)(nil).Serve))
+}
+
+// Stop mocks base method.
+func (m *MockGC) Stop() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Stop")
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockGCMockRecorder) Stop() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockGC)(nil).Stop))
+}

scheduler/supervisor/mocks/host_mock.go

@@ -0,0 +1,74 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: d7y.io/dragonfly/v2/scheduler/supervisor (interfaces: HostManager)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	reflect "reflect"
+
+	supervisor "d7y.io/dragonfly/v2/scheduler/supervisor"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockHostManager is a mock of HostManager interface.
+type MockHostManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockHostManagerMockRecorder
+}
+
+// MockHostManagerMockRecorder is the mock recorder for MockHostManager.
+type MockHostManagerMockRecorder struct {
+	mock *MockHostManager
+}
+
+// NewMockHostManager creates a new mock instance.
+func NewMockHostManager(ctrl *gomock.Controller) *MockHostManager {
+	mock := &MockHostManager{ctrl: ctrl}
+	mock.recorder = &MockHostManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockHostManager) EXPECT() *MockHostManagerMockRecorder {
+	return m.recorder
+}
+
+// Add mocks base method.
+func (m *MockHostManager) Add(arg0 *supervisor.Host) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Add", arg0)
+}
+
+// Add indicates an expected call of Add.
+func (mr *MockHostManagerMockRecorder) Add(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Add", reflect.TypeOf((*MockHostManager)(nil).Add), arg0)
+}
+
+// Delete mocks base method.
+func (m *MockHostManager) Delete(arg0 string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Delete", arg0)
+}
+
+// Delete indicates an expected call of Delete.
+func (mr *MockHostManagerMockRecorder) Delete(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockHostManager)(nil).Delete), arg0)
+}
+
+// Get mocks base method.
+func (m *MockHostManager) Get(arg0 string) (*supervisor.Host, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Get", arg0)
+	ret0, _ := ret[0].(*supervisor.Host)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// Get indicates an expected call of Get.
+func (mr *MockHostManagerMockRecorder) Get(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockHostManager)(nil).Get), arg0)
+}

scheduler/supervisor/mocks/peer_mock.go

@@ -0,0 +1,103 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: d7y.io/dragonfly/v2/scheduler/supervisor (interfaces: PeerManager)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	reflect "reflect"
+	sync "sync"
+
+	supervisor "d7y.io/dragonfly/v2/scheduler/supervisor"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockPeerManager is a mock of PeerManager interface.
+type MockPeerManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockPeerManagerMockRecorder
+}
+
+// MockPeerManagerMockRecorder is the mock recorder for MockPeerManager.
+type MockPeerManagerMockRecorder struct {
+	mock *MockPeerManager
+}
+
+// NewMockPeerManager creates a new mock instance.
+func NewMockPeerManager(ctrl *gomock.Controller) *MockPeerManager {
+	mock := &MockPeerManager{ctrl: ctrl}
+	mock.recorder = &MockPeerManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockPeerManager) EXPECT() *MockPeerManagerMockRecorder {
+	return m.recorder
+}
+
+// Add mocks base method.
+func (m *MockPeerManager) Add(arg0 *supervisor.Peer) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Add", arg0)
+}
+
+// Add indicates an expected call of Add.
+func (mr *MockPeerManagerMockRecorder) Add(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Add", reflect.TypeOf((*MockPeerManager)(nil).Add), arg0)
+}
+
+// Delete mocks base method.
+func (m *MockPeerManager) Delete(arg0 string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Delete", arg0)
+}
+
+// Delete indicates an expected call of Delete.
+func (mr *MockPeerManagerMockRecorder) Delete(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockPeerManager)(nil).Delete), arg0)
+}
+
+// Get mocks base method.
+func (m *MockPeerManager) Get(arg0 string) (*supervisor.Peer, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Get", arg0)
+	ret0, _ := ret[0].(*supervisor.Peer)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// Get indicates an expected call of Get.
+func (mr *MockPeerManagerMockRecorder) Get(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockPeerManager)(nil).Get), arg0)
+}
+
+// GetPeers mocks base method.
+func (m *MockPeerManager) GetPeers() *sync.Map {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPeers")
+	ret0, _ := ret[0].(*sync.Map)
+	return ret0
+}
+
+// GetPeers indicates an expected call of GetPeers.
+func (mr *MockPeerManagerMockRecorder) GetPeers() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPeers", reflect.TypeOf((*MockPeerManager)(nil).GetPeers))
+}
+
+// GetPeersByTask mocks base method.
+func (m *MockPeerManager) GetPeersByTask(arg0 string) []*supervisor.Peer {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPeersByTask", arg0)
+	ret0, _ := ret[0].([]*supervisor.Peer)
+	return ret0
+}
+
+// GetPeersByTask indicates an expected call of GetPeersByTask.
+func (mr *MockPeerManagerMockRecorder) GetPeersByTask(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPeersByTask", reflect.TypeOf((*MockPeerManager)(nil).GetPeersByTask), arg0)
+}

scheduler/supervisor/mocks/task_mock.go

@@ -0,0 +1,89 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: d7y.io/dragonfly/v2/scheduler/supervisor (interfaces: TaskManager)
+
+// Package mocks is a generated GoMock package.
+package mocks
+
+import (
+	reflect "reflect"
+
+	supervisor "d7y.io/dragonfly/v2/scheduler/supervisor"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockTaskManager is a mock of TaskManager interface.
+type MockTaskManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockTaskManagerMockRecorder
+}
+
+// MockTaskManagerMockRecorder is the mock recorder for MockTaskManager.
+type MockTaskManagerMockRecorder struct {
+	mock *MockTaskManager
+}
+
+// NewMockTaskManager creates a new mock instance.
+func NewMockTaskManager(ctrl *gomock.Controller) *MockTaskManager {
+	mock := &MockTaskManager{ctrl: ctrl}
+	mock.recorder = &MockTaskManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockTaskManager) EXPECT() *MockTaskManagerMockRecorder {
+	return m.recorder
+}
+
+// Add mocks base method.
+func (m *MockTaskManager) Add(arg0 *supervisor.Task) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Add", arg0)
+}
+
+// Add indicates an expected call of Add.
+func (mr *MockTaskManagerMockRecorder) Add(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Add", reflect.TypeOf((*MockTaskManager)(nil).Add), arg0)
+}
+
+// Delete mocks base method.
+func (m *MockTaskManager) Delete(arg0 string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Delete", arg0)
+}
+
+// Delete indicates an expected call of Delete.
+func (mr *MockTaskManagerMockRecorder) Delete(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockTaskManager)(nil).Delete), arg0)
+}
+
+// Get mocks base method.
+func (m *MockTaskManager) Get(arg0 string) (*supervisor.Task, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Get", arg0)
+	ret0, _ := ret[0].(*supervisor.Task)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// Get indicates an expected call of Get.
+func (mr *MockTaskManagerMockRecorder) Get(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockTaskManager)(nil).Get), arg0)
+}
+
+// GetOrAdd mocks base method.
+func (m *MockTaskManager) GetOrAdd(arg0 *supervisor.Task) (*supervisor.Task, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetOrAdd", arg0)
+	ret0, _ := ret[0].(*supervisor.Task)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// GetOrAdd indicates an expected call of GetOrAdd.
+func (mr *MockTaskManagerMockRecorder) GetOrAdd(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrAdd", reflect.TypeOf((*MockTaskManager)(nil).GetOrAdd), arg0)
+}

scheduler/supervisor/peer.go

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+//go:generate mockgen -destination ./mocks/peer_mock.go -package mocks d7y.io/dragonfly/v2/scheduler/supervisor PeerManager
 
 package supervisor
 
@@ -289,7 +290,7 @@ func isDescendant(ancestor, offspring *Peer) bool {
 		parent, ok := node.GetParent()
 		if !ok {
 			return false
-		} else if node.ID == ancestor.ID {
+		} else if parent.ID == ancestor.ID {
 			return true
 		}
 		node = parent

scheduler/supervisor/peer_test.go

@@ -0,0 +1,676 @@
+/*
+ *     Copyright 2020 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package supervisor_test
+
+import (
+	"strconv"
+	"testing"
+
+	"d7y.io/dragonfly/v2/scheduler/config"
+	"d7y.io/dragonfly/v2/scheduler/supervisor"
+	"d7y.io/dragonfly/v2/scheduler/supervisor/mocks"
+
+	"github.com/golang/mock/gomock"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPeer_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		id     string
+		expect func(t *testing.T, peer *supervisor.Peer)
+	}{
+		{
+			name: "create by normal config",
+			id:   "normal",
+			expect: func(t *testing.T, peer *supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Equal("normal", peer.ID)
+			},
+		},
+		{
+			name: "create by special symbols",
+			id:   "#@+:\b\t\\\"☹ ☺ ☻ (✿◠‿◠)",
+			expect: func(t *testing.T, peer *supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Equal("#@+:\b\t\\\"☹ ☺ ☻ (✿◠‿◠)", peer.ID)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			host := mockAHost("host")
+			peer := supervisor.NewPeer(tc.id, task, host)
+			tc.expect(t, peer)
+		})
+	}
+}
+
+func TestPeer_Tree(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		tree   map[int]int
+		answer []int
+		expect func(t *testing.T, peers []*supervisor.Peer, number int, answer []int)
+	}{
+		{
+			name:   "test ID of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{0, 1, 2, 3, 4, 5},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				for i := 0; i < number; i++ {
+					assert.Equal(strconv.Itoa(answer[i]), peers[i].ID)
+				}
+			},
+		},
+		{
+			name:   "test TreeNodeCount of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{6, 3, 2, 1, 1, 1},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				for i := 0; i < number; i++ {
+					assert.Equal(answer[i], peers[i].GetTreeNodeCount())
+				}
+			},
+		},
+		{
+			name:   "test TreeDepth of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{1, 2, 2, 3, 3, 3},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				for i := 0; i < number; i++ {
+					assert.Equal(answer[i], peers[i].GetTreeDepth())
+				}
+			},
+		},
+		{
+			name:   "test Root of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{0, 0, 0, 0, 0, 0},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				for i := 0; i < number; i++ {
+					assert.Equal(strconv.Itoa(answer[i]), peers[i].GetRoot().ID)
+				}
+			},
+		},
+		{
+			name:   "test Parent of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{-1, 0, 0, 1, 1, 2},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				for i := 0; i < number; i++ {
+					parent, success := peers[i].GetParent()
+					if answer[i] < 0 {
+						assert.Equal((*supervisor.Peer)(nil), parent)
+						assert.False(success)
+					} else {
+						assert.Equal(strconv.Itoa(answer[i]), parent.ID)
+						assert.True(success)
+					}
+				}
+			},
+		},
+		{
+			name:   "test Ancestor of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				assert.False(peers[0].IsAncestor(peers[0]))
+				assert.False(peers[0].IsAncestor(nil))
+
+				assert.True(peers[0].IsAncestor(peers[5]))
+				assert.False(peers[5].IsAncestor(peers[0]))
+
+				assert.True(peers[1].IsAncestor(peers[4]))
+				assert.False(peers[4].IsAncestor(peers[1]))
+			},
+		},
+		{
+			name:   "test Descendant of tree structure",
+			number: 6,
+			tree:   map[int]int{1: 0, 2: 0, 3: 1, 4: 1, 5: 2},
+			answer: []int{},
+			expect: func(t *testing.T, peers []*supervisor.Peer, number int, answer []int) {
+				assert := assert.New(t)
+				assert.False(peers[0].IsDescendant(peers[0]))
+				assert.False(peers[0].IsDescendant(nil))
+
+				assert.False(peers[0].IsDescendant(peers[5]))
+				assert.True(peers[5].IsDescendant(peers[0]))
+
+				assert.False(peers[1].IsDescendant(peers[4]))
+				assert.True(peers[4].IsDescendant(peers[1]))
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			var peers []*supervisor.Peer
+			task := mockATask("task")
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				peer := mockAPeer(index, task)
+				if i > 0 {
+					peer.ReplaceParent(peers[tc.tree[i]])
+				}
+				peers = append(peers, peer)
+			}
+			tc.expect(t, peers, tc.number, tc.answer)
+		})
+	}
+}
+
+func TestPeer_Cost(t *testing.T) {
+	tests := []struct {
+		name          string
+		finishedCount []int32
+		cost          []int
+		expect        func(t *testing.T, peer *supervisor.Peer, cost []int)
+	}{
+		{
+			name:          "normal workflow",
+			finishedCount: []int32{2, 3, 4},
+			cost:          []int{3, 4, 5},
+			expect: func(t *testing.T, peer *supervisor.Peer, cost []int) {
+				assert := assert.New(t)
+
+				costFetch := peer.GetPieceCosts()
+				assert.ElementsMatch(costFetch, cost)
+
+				average, success := peer.GetPieceAverageCost()
+				assert.True(success)
+				assert.Equal(4, average)
+				finishedCountFetch, loadFetch := peer.GetSortKeys()
+				assert.Equal(4, finishedCountFetch)
+				assert.Equal(100, loadFetch)
+			},
+		},
+		{
+			name:          "no workflow will be neglected",
+			finishedCount: []int32{},
+			cost:          []int{},
+			expect: func(t *testing.T, peer *supervisor.Peer, cost []int) {
+				assert := assert.New(t)
+
+				average, success := peer.GetPieceAverageCost()
+				assert.False(success)
+				assert.Equal(0, average)
+			},
+		},
+		{
+			name: "long workflow will be clipped",
+			finishedCount: []int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+				11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22},
+			cost: []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+				11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22},
+			expect: func(t *testing.T, peer *supervisor.Peer, cost []int) {
+				assert := assert.New(t)
+
+				costFetch := peer.GetPieceCosts()
+				assert.ElementsMatch(costFetch, cost[2:])
+
+				average, success := peer.GetPieceAverageCost()
+				assert.True(success)
+				assert.Equal(12, average)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			peer := mockAPeer("peer", task)
+			for i := 0; i < len(tc.finishedCount); i++ {
+				peer.UpdateProgress(tc.finishedCount[i], tc.cost[i])
+			}
+			tc.expect(t, peer, tc.cost)
+		})
+	}
+}
+
+func TestPeer_Status(t *testing.T) {
+	tests := []struct {
+		name       string
+		status     supervisor.PeerStatus
+		statusName string
+		judgeArray []bool
+		expect     func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool)
+	}{
+		{
+			name:       "status Waiting",
+			status:     supervisor.PeerStatusWaiting,
+			statusName: "Waiting",
+			judgeArray: []bool{false, true, false, false, false, false},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Running",
+			status:     supervisor.PeerStatusRunning,
+			statusName: "Running",
+			judgeArray: []bool{true, false, false, false, false, false},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Zombie",
+			status:     supervisor.PeerStatusZombie,
+			statusName: "Zombie",
+			judgeArray: []bool{false, false, false, false, true, false},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Fail",
+			status:     supervisor.PeerStatusFail,
+			statusName: "Fail",
+			judgeArray: []bool{false, false, false, true, true, true},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Success",
+			status:     supervisor.PeerStatusSuccess,
+			statusName: "Success",
+			judgeArray: []bool{false, false, true, true, false, false},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "unknown",
+			status:     100,
+			statusName: "unknown",
+			judgeArray: []bool{false, false, false, false, false, false},
+			expect: func(t *testing.T, peer *supervisor.Peer, status supervisor.PeerStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.PeerStatus.String(status), statusName)
+				assert.Equal(peer.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					peer.IsRunning(), peer.IsWaiting(), peer.IsSuccess(),
+					peer.IsDone(), peer.IsBad(), peer.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			peer := mockAPeer("peer", task)
+			peer.SetStatus(tc.status)
+
+			tc.expect(t, peer, tc.status, tc.statusName, tc.judgeArray)
+		})
+	}
+}
+
+func TestPeerManager_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		config *config.GCConfig
+		mock   func(m *mocks.MockGCMockRecorder)
+		expect func(t *testing.T, peerManager supervisor.PeerManager, err error)
+	}{
+		{
+			name:   "create with default config",
+			config: config.New().Scheduler.GC,
+			mock: func(m *mocks.MockGCMockRecorder) {
+				m.Add(gomock.Any()).Return(nil).AnyTimes()
+			},
+			expect: func(t *testing.T, peerManager supervisor.PeerManager, err error) {
+				assert := assert.New(t)
+				assert.NotNil(peerManager)
+				assert.Nil(err)
+			},
+		},
+		{
+			name: "create with strange int",
+			config: &config.GCConfig{
+				PeerGCInterval: 1,
+				TaskGCInterval: 1 >> 69,
+				PeerTTL:        1 << 62,
+				PeerTTI:        1,
+				TaskTTL:        1,
+				TaskTTI:        1,
+			},
+			mock: func(m *mocks.MockGCMockRecorder) {
+				m.Add(gomock.Any()).Return(nil).AnyTimes()
+			},
+			expect: func(t *testing.T, peerManager supervisor.PeerManager, err error) {
+				assert := assert.New(t)
+				assert.NotNil(peerManager)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "gc failed",
+			config: config.New().Scheduler.GC,
+			mock: func(m *mocks.MockGCMockRecorder) {
+				m.Add(gomock.Any()).Return(errors.New("mockError")).AnyTimes()
+			},
+			expect: func(t *testing.T, peerManager supervisor.PeerManager, err error) {
+				assert := assert.New(t)
+				assert.Nil(peerManager)
+				assert.Error(err)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockHostManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			tc.mock(mockGC.EXPECT())
+
+			peerManager, err := supervisor.NewPeerManager(tc.config, mockGC, mockHostManager)
+			tc.expect(t, peerManager, err)
+		})
+	}
+}
+
+func TestPeerManager_GetPeer(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		fetch  int
+		expect func(t *testing.T, peer *supervisor.Peer, success bool, err error)
+	}{
+		{
+			name:   "fetch first peer",
+			number: 3,
+			fetch:  0,
+			expect: func(t *testing.T, peer *supervisor.Peer, success bool, err error) {
+				assert := assert.New(t)
+				assert.Equal("0", peer.ID)
+				assert.True(success)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "fetch last peer",
+			number: 3,
+			fetch:  2,
+			expect: func(t *testing.T, peer *supervisor.Peer, success bool, err error) {
+				assert := assert.New(t)
+				assert.Equal("2", peer.ID)
+				assert.True(success)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "fetch not exist peer",
+			number: 3,
+			fetch:  -1,
+			expect: func(t *testing.T, peer *supervisor.Peer, success bool, err error) {
+				assert := assert.New(t)
+				assert.Equal((*supervisor.Peer)(nil), peer)
+				assert.False(success)
+				assert.Nil(err)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockHostManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			peerManager, err := supervisor.NewPeerManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			task := mockATask("123")
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				peer := mockAPeer(index, task)
+				peerManager.Add(peer)
+			}
+			peer, success := peerManager.Get(strconv.Itoa(tc.fetch))
+			tc.expect(t, peer, success, err)
+		})
+	}
+}
+
+func TestPeerManager_Add(t *testing.T) {
+	tests := []struct {
+		name   string
+		ID     []int
+		expect func(t *testing.T, peer []*supervisor.Peer)
+	}{
+		{
+			name: "add seperative peers",
+			ID:   []int{1, 2, 3},
+			expect: func(t *testing.T, peers []*supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Len(peers, 3)
+			},
+		},
+		{
+			name: "add duplicate peers",
+			ID:   []int{1, 1, 1},
+			expect: func(t *testing.T, peers []*supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Len(peers, 1)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockHostManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			peerManager, _ := supervisor.NewPeerManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			task := mockATask("123")
+			for _, i := range tc.ID {
+				index := strconv.Itoa(i)
+				peer := mockAPeer(index, task)
+				peerManager.Add(peer)
+			}
+			peers := peerManager.GetPeersByTask("123")
+			tc.expect(t, peers)
+		})
+	}
+}
+
+func TestPeerManager_GetPeersByTask(t *testing.T) {
+	tests := []struct {
+		name   string
+		tasks  map[*supervisor.Task]int
+		fetch  string
+		expect func(t *testing.T, peer []*supervisor.Peer)
+	}{
+		{
+			name:  "peer for a task",
+			tasks: map[*supervisor.Task]int{mockATask("123"): 3},
+			fetch: "123",
+			expect: func(t *testing.T, peers []*supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Len(peers, 3)
+			},
+		},
+		{
+			name:  "one from two task",
+			tasks: map[*supervisor.Task]int{mockATask("123"): 2, mockATask("456"): 3},
+			fetch: "123",
+			expect: func(t *testing.T, peers []*supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Len(peers, 2)
+			},
+		},
+		{
+			name:  "no peer for a task",
+			tasks: map[*supervisor.Task]int{mockATask("123"): 1},
+			fetch: "456",
+			expect: func(t *testing.T, peers []*supervisor.Peer) {
+				assert := assert.New(t)
+				assert.Len(peers, 0)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockHostManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			peerManager, _ := supervisor.NewPeerManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			nowAt := 0
+			for task, num := range tc.tasks {
+				for i := nowAt; i < nowAt+num; i++ {
+					index := strconv.Itoa(i)
+					peer := mockAPeer(index, task)
+					t.Log(i, index, nowAt, peer.ID, num)
+					peerManager.Add(peer)
+				}
+				nowAt += num
+			}
+			peers := peerManager.GetPeersByTask(tc.fetch)
+			tc.expect(t, peers)
+		})
+	}
+}
+
+func TestPeerManager_Delete(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		delete int
+		fetch  int
+		expect func(t *testing.T, peer *supervisor.Peer, success bool)
+	}{
+		{
+			name:   "delete exist peer",
+			number: 1,
+			delete: 0,
+			fetch:  0,
+			expect: func(t *testing.T, peer *supervisor.Peer, success bool) {
+				assert := assert.New(t)
+				assert.Nil(peer)
+				assert.False(success)
+			},
+		},
+		{
+			name:   "delete not exist peer",
+			number: 1,
+			delete: 100,
+			fetch:  0,
+			expect: func(t *testing.T, peer *supervisor.Peer, success bool) {
+				assert := assert.New(t)
+				assert.NotNil(peer)
+				assert.True(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockHostManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			peerManager, _ := supervisor.NewPeerManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			task := mockATask("123")
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				peer := mockAPeer(index, task)
+				peerManager.Add(peer)
+			}
+			peerManager.Delete(strconv.Itoa(tc.delete))
+			peer, success := peerManager.Get(strconv.Itoa(tc.fetch))
+
+			tc.expect(t, peer, success)
+		})
+	}
+}
+
+func mockAPeer(ID string, task *supervisor.Task) *supervisor.Peer {
+	host := supervisor.NewClientHost(ID, "127.0.0.1", "Client", 8080, 8081, "", "", "")
+	return supervisor.NewPeer(ID, task, host)
+}

scheduler/supervisor/task.go

@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+//go:generate mockgen -destination ./mocks/task_mock.go -package mocks d7y.io/dragonfly/v2/scheduler/supervisor TaskManager
 
 package supervisor
 

scheduler/supervisor/task_test.go

@@ -0,0 +1,576 @@
+/*
+ *     Copyright 2020 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package supervisor_test
+
+import (
+	"strconv"
+	"testing"
+
+	"d7y.io/dragonfly/v2/pkg/rpc/base"
+	"d7y.io/dragonfly/v2/scheduler/config"
+	"d7y.io/dragonfly/v2/scheduler/supervisor"
+	"d7y.io/dragonfly/v2/scheduler/supervisor/mocks"
+	"github.com/golang/mock/gomock"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTask_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		task   *supervisor.Task
+		expect func(t *testing.T, task *supervisor.Task)
+	}{
+		{
+			name: "create by normal config",
+			task: supervisor.NewTask("main", "127.0.0.1", &base.UrlMeta{}),
+			expect: func(t *testing.T, task *supervisor.Task) {
+				assert := assert.New(t)
+				assert.Equal("main", task.ID)
+			},
+		},
+		{
+			name: "create by special symbol",
+			task: supervisor.NewTask("\x07\b%$!！\x7F✌ (>‿<)✌", "d7y.io/dragonfly", &base.UrlMeta{Tag: "d7y-test"}),
+			expect: func(t *testing.T, task *supervisor.Task) {
+				assert := assert.New(t)
+				assert.Equal("\x07\b%$!！\x7F✌ (>‿<)✌", task.ID)
+			},
+		},
+		{
+			name: "create by http url",
+			task: supervisor.NewTask("task", "http://370.moe/", &base.UrlMeta{}),
+			expect: func(t *testing.T, task *supervisor.Task) {
+				assert := assert.New(t)
+				assert.Equal("task", task.ID)
+			},
+		},
+		{
+			name: "create by normal config",
+			task: supervisor.NewTask("task", "android://370.moe", &base.UrlMeta{}),
+			expect: func(t *testing.T, task *supervisor.Task) {
+				assert := assert.New(t)
+				assert.Equal("task", task.ID)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tc.expect(t, tc.task)
+		})
+	}
+}
+
+func TestTask_Status(t *testing.T) {
+	tests := []struct {
+		name       string
+		status     supervisor.TaskStatus
+		statusName string
+		judgeArray []bool
+		expect     func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool)
+	}{
+		{
+			name:       "status Waiting",
+			status:     supervisor.TaskStatusWaiting,
+			statusName: "Waiting",
+			judgeArray: []bool{false, false, true, false, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Running",
+			status:     supervisor.TaskStatusRunning,
+			statusName: "Running",
+			judgeArray: []bool{false, false, false, true, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status seeding",
+			status:     supervisor.TaskStatusSeeding,
+			statusName: "Seeding",
+			judgeArray: []bool{false, true, false, true, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status success",
+			status:     supervisor.TaskStatusSuccess,
+			statusName: "Success",
+			judgeArray: []bool{true, true, false, true, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status zombie",
+			status:     supervisor.TaskStatusZombie,
+			statusName: "Zombie",
+			judgeArray: []bool{false, false, false, false, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "status Fail",
+			status:     supervisor.TaskStatusFail,
+			statusName: "Fail",
+			judgeArray: []bool{false, false, false, false, true},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+		{
+			name:       "unknown",
+			status:     100,
+			statusName: "unknown",
+			judgeArray: []bool{false, false, false, false, false},
+			expect: func(t *testing.T, task *supervisor.Task, status supervisor.TaskStatus, statusName string, judgeArray []bool) {
+				assert := assert.New(t)
+				assert.Equal(supervisor.TaskStatus.String(status), statusName)
+				assert.Equal(task.GetStatus(), status)
+
+				statutusJudgeArray := []bool{
+					task.IsSuccess(), task.CanSchedule(),
+					task.IsWaiting(), task.IsHealth(), task.IsFail(),
+				}
+				assert.Equal(statutusJudgeArray, judgeArray)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			task.SetStatus(tc.status)
+			tc.expect(t, task, tc.status, tc.statusName, tc.judgeArray)
+		})
+	}
+}
+
+func TestTask_BackToSourcePeer(t *testing.T) {
+	tests := []struct {
+		name          string
+		initialWeight int32
+		add           []string
+		expect        func(t *testing.T, task *supervisor.Task, add []string)
+	}{
+		{
+			name:          "able to backsource",
+			initialWeight: 4,
+			add:           []string{"0", "1", "2"},
+			expect: func(t *testing.T, task *supervisor.Task, add []string) {
+				assert := assert.New(t)
+				assert.EqualValues(task.BackToSourceWeight.Load(), 1)
+				assert.True(task.CanBackToSource())
+				assert.ElementsMatch(task.GetBackToSourcePeers(), add)
+				for _, ID := range add {
+					contain := task.ContainsBackToSourcePeer(ID)
+					assert.True(contain)
+				}
+
+			},
+		},
+		{
+			name:          "unable to backsource",
+			initialWeight: -1,
+			add:           []string{},
+			expect: func(t *testing.T, task *supervisor.Task, add []string) {
+				assert := assert.New(t)
+				assert.EqualValues(task.BackToSourceWeight.Load(), -1)
+				assert.False(task.CanBackToSource())
+
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			task.BackToSourceWeight.Store(tc.initialWeight)
+			for _, ID := range tc.add {
+				task.AddBackToSourcePeer(ID)
+			}
+			tc.expect(t, task, tc.add)
+		})
+	}
+}
+
+func TestTask_Pick(t *testing.T) {
+	tests := []struct {
+		name    string
+		number  int
+		pick    func(peer *supervisor.Peer) bool
+		reverse bool
+		limit   int
+		answer  []string
+	}{
+		{
+			name:   "pick three odd",
+			number: 10,
+			pick: func(peer *supervisor.Peer) bool {
+				id, _ := strconv.Atoi(peer.ID)
+				return id%2 != 0
+			},
+			reverse: false,
+			limit:   3,
+			answer:  []string{"1", "3", "5"},
+		},
+		{
+			name:   "pick 100 odd",
+			number: 10,
+			pick: func(peer *supervisor.Peer) bool {
+				id, _ := strconv.Atoi(peer.ID)
+				return id%2 != 0
+			},
+			reverse: true,
+			limit:   3,
+			answer:  []string{"5", "7", "9"},
+		},
+		{
+			name:   "pick all odd",
+			number: 10,
+			pick: func(peer *supervisor.Peer) bool {
+				id, _ := strconv.Atoi(peer.ID)
+				return id%2 != 0
+			},
+			reverse: false,
+			limit:   100,
+			answer:  []string{"1", "3", "5", "7", "9"},
+		},
+		{
+			name:   "pick all",
+			number: 10,
+			pick: func(peer *supervisor.Peer) bool {
+				return true
+			},
+			reverse: false,
+			limit:   100,
+			answer:  []string{"0", "1", "2", "3", "4", "5", "6", "7", "8", "9"},
+		},
+		{
+			name:   "pick nil",
+			number: 10,
+			pick: func(peer *supervisor.Peer) bool {
+				return false
+			},
+			reverse: false,
+			limit:   100,
+			answer:  []string{},
+		},
+		{
+			name:    "invalid pickFn",
+			number:  10,
+			pick:    (func(peer *supervisor.Peer) bool)(nil),
+			reverse: false,
+			limit:   100,
+			answer:  []string{},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			task := mockATask("task")
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				peer := mockAPeer(index, task)
+				peer.UpdateProgress((int32)(i), i)
+				task.AddPeer(peer)
+			}
+			var peers []*supervisor.Peer
+			if tc.reverse {
+				peers = task.PickReverse(tc.limit, tc.pick)
+			} else {
+				peers = task.Pick(tc.limit, tc.pick)
+			}
+			var peerIDs []string
+			for _, peer := range peers {
+				peerIDs = append(peerIDs, peer.ID)
+			}
+			assert := assert.New(t)
+			assert.ElementsMatch(peerIDs, tc.answer)
+		})
+	}
+}
+
+func TestTaskManager_New(t *testing.T) {
+	tests := []struct {
+		name   string
+		config *config.GCConfig
+		mock   func(m *mocks.MockGCMockRecorder)
+		expect func(t *testing.T, taskManager supervisor.TaskManager, err error)
+	}{
+		{
+			name:   "simple create",
+			config: config.New().Scheduler.GC,
+			mock: func(m *mocks.MockGCMockRecorder) {
+				m.Add(gomock.Any()).Return(nil).AnyTimes()
+			},
+			expect: func(t *testing.T, taskManager supervisor.TaskManager, err error) {
+				assert := assert.New(t)
+				assert.NotNil(taskManager)
+				assert.Nil(err)
+			},
+		},
+		{
+			name:   "gc failed",
+			config: config.New().Scheduler.GC,
+			mock: func(m *mocks.MockGCMockRecorder) {
+				m.Add(gomock.Any()).Return(errors.New("mockError")).AnyTimes()
+			},
+			expect: func(t *testing.T, taskManager supervisor.TaskManager, err error) {
+				assert := assert.New(t)
+				assert.Nil(taskManager)
+				assert.Error(err)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockAPeerManager := mocks.NewMockPeerManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			tc.mock(mockGC.EXPECT())
+
+			taskManager, err := supervisor.NewTaskManager(tc.config, mockGC, mockAPeerManager)
+			tc.expect(t, taskManager, err)
+		})
+	}
+}
+
+func TestTaskManager_Get(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		fetch  int
+		expect func(t *testing.T, task *supervisor.Task, success bool)
+	}{
+		{
+			name:   "fetch first task",
+			number: 3,
+			fetch:  0,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Equal("0", task.ID)
+				assert.True(success)
+			},
+		},
+		{
+			name:   "fetch last task",
+			number: 3,
+			fetch:  2,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Equal("2", task.ID)
+				assert.True(success)
+			},
+		},
+		{
+			name:   "fetch not exist task",
+			number: 3,
+			fetch:  -1,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Equal((*supervisor.Task)(nil), task)
+				assert.False(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockPeerManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			taskManager, _ := supervisor.NewTaskManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				task := mockATask(index)
+				taskManager.Add(task)
+			}
+			task, success := taskManager.Get(strconv.Itoa(tc.fetch))
+			tc.expect(t, task, success)
+		})
+	}
+}
+
+func TestTaskManager_GetOrAdd(t *testing.T) {
+	tests := []struct {
+		name   string
+		create int
+		add    int
+		expect func(t *testing.T, task *supervisor.Task, success bool)
+	}{
+		{
+			name:   "get exist task",
+			create: 3,
+			add:    0,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Equal("2", task.ID)
+				assert.False(success)
+			},
+		},
+		{
+			name:   "add not exist task",
+			create: 3,
+			add:    3,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Equal("2", task.ID)
+				assert.True(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockHostManager := mocks.NewMockPeerManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			taskManager, _ := supervisor.NewTaskManager(cfg.Scheduler.GC, mockGC, mockHostManager)
+			var tasks []*supervisor.Task
+			for i := 0; i < tc.create; i++ {
+				index := strconv.Itoa(i)
+				task := mockATask(index)
+				tasks = append(tasks, task)
+			}
+			for i := 0; i < tc.add; i++ {
+				taskManager.Add(tasks[i])
+			}
+			task, success := taskManager.GetOrAdd(tasks[len(tasks)-1])
+			tc.expect(t, task, success)
+		})
+	}
+}
+
+func TestTaskManager_Delete(t *testing.T) {
+	tests := []struct {
+		name   string
+		number int
+		delete int
+		fetch  int
+		expect func(t *testing.T, task *supervisor.Task, success bool)
+	}{
+		{
+			name:   "delete exist task",
+			number: 1,
+			delete: 0,
+			fetch:  0,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.Nil(task)
+				assert.False(success)
+			},
+		},
+		{
+			name:   "delete not exist task",
+			number: 1,
+			delete: 100,
+			fetch:  0,
+			expect: func(t *testing.T, task *supervisor.Task, success bool) {
+				assert := assert.New(t)
+				assert.NotNil(task)
+				assert.True(success)
+			},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctl := gomock.NewController(t)
+			defer ctl.Finish()
+			mockAPeerManager := mocks.NewMockPeerManager(ctl)
+			mockGC := mocks.NewMockGC(ctl)
+			mockGC.EXPECT().Add(gomock.Any()).Return(nil).AnyTimes()
+
+			cfg := config.New()
+			taskManager, _ := supervisor.NewTaskManager(cfg.Scheduler.GC, mockGC, mockAPeerManager)
+			for i := 0; i < tc.number; i++ {
+				index := strconv.Itoa(i)
+				task := mockATask(index)
+				taskManager.Add(task)
+			}
+			taskManager.Delete(strconv.Itoa(tc.delete))
+			task, success := taskManager.Get(strconv.Itoa(tc.fetch))
+
+			tc.expect(t, task, success)
+		})
+	}
+}
+
+func mockATask(ID string) *supervisor.Task {
+	urlMeta := &base.UrlMeta{
+		Tag: "d7y-test",
+	}
+	return supervisor.NewTask(ID, "d7y.io/dragonfly", urlMeta)
+}