577 lines
16 KiB
Go
577 lines
16 KiB
Go
/*
|
||
* 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)
|
||
}
|