github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/mocks/mocks.go (about)

     1  package mocks
     2  
     3  import (
     4  	model "github.com/cloudreve/Cloudreve/v3/models"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
     6  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
     7  	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
     8  	"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
     9  	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
    10  	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
    11  	"github.com/cloudreve/Cloudreve/v3/pkg/task"
    12  	testMock "github.com/stretchr/testify/mock"
    13  )
    14  
    15  type NodePoolMock struct {
    16  	testMock.Mock
    17  }
    18  
    19  func (n NodePoolMock) BalanceNodeByFeature(feature string, lb balancer.Balancer) (error, cluster.Node) {
    20  	args := n.Called(feature, lb)
    21  	return args.Error(0), args.Get(1).(cluster.Node)
    22  }
    23  
    24  func (n NodePoolMock) GetNodeByID(id uint) cluster.Node {
    25  	args := n.Called(id)
    26  	if res, ok := args.Get(0).(cluster.Node); ok {
    27  		return res
    28  	}
    29  
    30  	return nil
    31  }
    32  
    33  func (n NodePoolMock) Add(node *model.Node) {
    34  	n.Called(node)
    35  }
    36  
    37  func (n NodePoolMock) Delete(id uint) {
    38  	n.Called(id)
    39  }
    40  
    41  type NodeMock struct {
    42  	testMock.Mock
    43  }
    44  
    45  func (n NodeMock) Init(node *model.Node) {
    46  	n.Called(node)
    47  }
    48  
    49  func (n NodeMock) IsFeatureEnabled(feature string) bool {
    50  	args := n.Called(feature)
    51  	return args.Bool(0)
    52  }
    53  
    54  func (n NodeMock) SubscribeStatusChange(callback func(isActive bool, id uint)) {
    55  	n.Called(callback)
    56  }
    57  
    58  func (n NodeMock) Ping(req *serializer.NodePingReq) (*serializer.NodePingResp, error) {
    59  	args := n.Called(req)
    60  	return args.Get(0).(*serializer.NodePingResp), args.Error(1)
    61  }
    62  
    63  func (n NodeMock) IsActive() bool {
    64  	args := n.Called()
    65  	return args.Bool(0)
    66  }
    67  
    68  func (n NodeMock) GetAria2Instance() common.Aria2 {
    69  	args := n.Called()
    70  	return args.Get(0).(common.Aria2)
    71  }
    72  
    73  func (n NodeMock) ID() uint {
    74  	args := n.Called()
    75  	return args.Get(0).(uint)
    76  }
    77  
    78  func (n NodeMock) Kill() {
    79  	n.Called()
    80  }
    81  
    82  func (n NodeMock) IsMater() bool {
    83  	args := n.Called()
    84  	return args.Bool(0)
    85  }
    86  
    87  func (n NodeMock) MasterAuthInstance() auth.Auth {
    88  	args := n.Called()
    89  	return args.Get(0).(auth.Auth)
    90  }
    91  
    92  func (n NodeMock) SlaveAuthInstance() auth.Auth {
    93  	args := n.Called()
    94  	return args.Get(0).(auth.Auth)
    95  }
    96  
    97  func (n NodeMock) DBModel() *model.Node {
    98  	args := n.Called()
    99  	return args.Get(0).(*model.Node)
   100  }
   101  
   102  type Aria2Mock struct {
   103  	testMock.Mock
   104  }
   105  
   106  func (a Aria2Mock) Init() error {
   107  	args := a.Called()
   108  	return args.Error(0)
   109  }
   110  
   111  func (a Aria2Mock) CreateTask(task *model.Download, options map[string]interface{}) (string, error) {
   112  	args := a.Called(task, options)
   113  	return args.String(0), args.Error(1)
   114  }
   115  
   116  func (a Aria2Mock) Status(task *model.Download) (rpc.StatusInfo, error) {
   117  	args := a.Called(task)
   118  	return args.Get(0).(rpc.StatusInfo), args.Error(1)
   119  }
   120  
   121  func (a Aria2Mock) Cancel(task *model.Download) error {
   122  	args := a.Called(task)
   123  	return args.Error(0)
   124  }
   125  
   126  func (a Aria2Mock) Select(task *model.Download, files []int) error {
   127  	args := a.Called(task, files)
   128  	return args.Error(0)
   129  }
   130  
   131  func (a Aria2Mock) GetConfig() model.Aria2Option {
   132  	args := a.Called()
   133  	return args.Get(0).(model.Aria2Option)
   134  }
   135  
   136  func (a Aria2Mock) DeleteTempFile(download *model.Download) error {
   137  	args := a.Called(download)
   138  	return args.Error(0)
   139  }
   140  
   141  type TaskPoolMock struct {
   142  	testMock.Mock
   143  }
   144  
   145  func (t TaskPoolMock) Add(num int) {
   146  	t.Called(num)
   147  }
   148  
   149  func (t TaskPoolMock) Submit(job task.Job) {
   150  	t.Called(job)
   151  }