mirror of
https://github.com/cloudreve/Cloudreve.git
synced 2024-04-21 12:31:40 +00:00
Compare commits
36
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
36bbf4fc21 | ||
|
|
a910a20457 | ||
|
|
c4fc4a803d | ||
|
|
c31c77a089 | ||
|
|
6b15cae0b5 | ||
|
|
84d81f201f | ||
|
|
af4d9767c2 | ||
|
|
45597adcd3 | ||
|
|
762f0f9c68 | ||
|
|
c5074df1c7 | ||
|
|
7ea72cf364 | ||
|
|
4eb7525c51 | ||
|
|
3948ee7f3a | ||
|
|
865a801fa8 | ||
|
|
05941616df | ||
|
|
51b1e5b854 | ||
|
|
4dbe867020 | ||
|
|
8c8ad3e149 | ||
|
|
fce38209bc | ||
|
|
700e13384e | ||
|
|
7fd984f95d | ||
|
|
9fc08292a0 | ||
|
|
8c5445a26d | ||
|
|
96b84bb5e5 | ||
|
|
9056ef9171 | ||
|
|
532bff820a | ||
|
|
fcd9eddc54 | ||
|
|
6c9967b120 | ||
|
|
416f4c1dd2 | ||
|
|
f0089045d7 | ||
|
|
4b88eacb6a | ||
|
|
54ed7e43ca | ||
|
|
4d7b8685b9 | ||
|
|
eeee43d569 | ||
|
|
3064ed60f3 | ||
|
|
e41ec9defa |
@@ -1,8 +1,6 @@
|
||||
name: Build
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ master ]
|
||||
on: workflow_dispatch
|
||||
|
||||
jobs:
|
||||
build:
|
||||
|
||||
+11
-4
@@ -5,6 +5,9 @@ COPY ./assets /assets
|
||||
|
||||
WORKDIR /assets
|
||||
|
||||
# If encountered problems like JavaScript heap out of memory, please uncomment the following options
|
||||
ENV NODE_OPTIONS --max_old_space_size=4096
|
||||
|
||||
# yarn repo connection is unstable, adjust the network timeout to 10 min.
|
||||
RUN set -ex \
|
||||
&& yarn install --network-timeout 600000 \
|
||||
@@ -42,14 +45,18 @@ ARG TZ="Asia/Shanghai"
|
||||
|
||||
ENV TZ ${TZ}
|
||||
|
||||
COPY --from=be-builder /go/bin/cloudreve /cloudreve/cloudreve
|
||||
COPY --from=be-builder /go/bin/Cloudreve /cloudreve/cloudreve
|
||||
COPY docker-bootstrap.sh /cloudreve/bootstrap.sh
|
||||
|
||||
RUN apk upgrade \
|
||||
&& apk add bash tzdata \
|
||||
&& apk add bash tzdata aria2 \
|
||||
&& ln -s /cloudreve/cloudreve /usr/bin/cloudreve \
|
||||
&& ln -sf /usr/share/zoneinfo/${TZ} /etc/localtime \
|
||||
&& echo ${TZ} > /etc/timezone \
|
||||
&& rm -rf /var/cache/apk/*
|
||||
&& rm -rf /var/cache/apk/* \
|
||||
&& mkdir /etc/cloudreve \
|
||||
&& ln -s /etc/cloudreve/cloureve.db /cloudreve/cloudreve.db \
|
||||
&& ln -s /etc/cloudreve/conf.ini /cloudreve/conf.ini
|
||||
|
||||
# cloudreve use tcp 5212 port by default
|
||||
EXPOSE 5212/tcp
|
||||
@@ -62,4 +69,4 @@ VOLUME /etc/cloudreve
|
||||
|
||||
VOLUME /data
|
||||
|
||||
ENTRYPOINT ["cloudreve"]
|
||||
ENTRYPOINT ["sh", "/cloudreve/bootstrap.sh"]
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
|
||||
* :cloud: 支持本机、从机、七牛、阿里云 OSS、腾讯云 COS、又拍云、OneDrive (包括世纪互联版) 作为存储端
|
||||
* :outbox_tray: 上传/下载 支持客户端直传,支持下载限速
|
||||
* 💾 可对接 Aria2 离线下载
|
||||
* 💾 可对接 Aria2 离线下载,可使用多个从机机点分担下载任务
|
||||
* 📚 在线 压缩/解压缩、多文件打包下载
|
||||
* 💻 覆盖全部存储策略的 WebDAV 协议支持
|
||||
* :zap: 拖拽上传、目录上传、流式上传处理
|
||||
|
||||
+1
-1
Submodule assets updated: efe2981761...6eda38c85a
+1
-1
@@ -15,7 +15,7 @@ func InitApplication() {
|
||||
fmt.Print(`
|
||||
___ _ _
|
||||
/ __\ | ___ _ _ __| |_ __ _____ _____
|
||||
/ / | |/ _ \| | | |/ _ | '__/ _ \ \ / / _ \
|
||||
/ / | |/ _ \| | | |/ _ | '__/ _ \ \ / / _ \
|
||||
/ /___| | (_) | |_| | (_| | | | __/\ V / __/
|
||||
\____/|_|\___/ \__,_|\__,_|_| \___| \_/ \___|
|
||||
|
||||
|
||||
+10
-3
@@ -2,6 +2,7 @@ package bootstrap
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/models/scripts"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
@@ -9,7 +10,7 @@ import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/crontab"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/email"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/slave"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
@@ -27,6 +28,12 @@ func Init(path string) {
|
||||
mode string
|
||||
factory func()
|
||||
}{
|
||||
{
|
||||
"both",
|
||||
func() {
|
||||
scripts.Init()
|
||||
},
|
||||
},
|
||||
{
|
||||
"both",
|
||||
func() {
|
||||
@@ -54,7 +61,7 @@ func Init(path string) {
|
||||
{
|
||||
"master",
|
||||
func() {
|
||||
aria2.Init(false)
|
||||
aria2.Init(false, cluster.Default, mq.GlobalMQ)
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -78,7 +85,7 @@ func Init(path string) {
|
||||
{
|
||||
"slave",
|
||||
func() {
|
||||
slave.Init()
|
||||
cluster.InitController()
|
||||
},
|
||||
},
|
||||
{
|
||||
|
||||
+2
-2
@@ -2,14 +2,14 @@ package bootstrap
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/cloudreve/Cloudreve/v3/models/scripts"
|
||||
"github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
)
|
||||
|
||||
func RunScript(name string) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
if err := scripts.RunDBScript(name, ctx); err != nil {
|
||||
if err := invoker.RunDBScript(name, ctx); err != nil {
|
||||
util.Log().Error("数据库脚本执行失败: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/sh
|
||||
GREEN='\033[0;32m'
|
||||
RESET='\033[0m'
|
||||
if [ ! -f /etc/cloudreve/aria2c.conf ]; then
|
||||
echo -e "[${GREEN}aria2c${RESET}] aria2c config not found. Generating..."
|
||||
secret=$(tr -dc A-Za-z0-9 </dev/urandom | head -c 13)
|
||||
echo -e "[${GREEN}aria2c${RESET}] Generated port: 6800, secret: $secret"
|
||||
cat <<EOF > /etc/cloudreve/aria2c.conf
|
||||
enable-rpc=true
|
||||
rpc-listen-port=6800
|
||||
rpc-secret=$secret
|
||||
EOF
|
||||
fi
|
||||
aria2c --conf-path /etc/cloudreve/aria2c.conf -D
|
||||
cloudreve
|
||||
@@ -38,7 +38,8 @@ require (
|
||||
github.com/tencentcloud/tencentcloud-sdk-go v3.0.125+incompatible
|
||||
github.com/tencentyun/cos-go-sdk-v5 v0.0.0-20200120023323-87ff3bc489ac
|
||||
github.com/upyun/go-sdk v2.1.0+incompatible
|
||||
golang.org/x/text v0.3.2
|
||||
golang.org/x/image v0.0.0-20211028202545-6944b10bf410
|
||||
golang.org/x/text v0.3.6
|
||||
gopkg.in/alexcesaro/quotedprintable.v3 v3.0.0-20150716171945-2caba252f4dc // indirect
|
||||
gopkg.in/go-playground/validator.v9 v9.29.1
|
||||
gopkg.in/ini.v1 v1.51.0 // indirect
|
||||
|
||||
@@ -12,6 +12,7 @@ github.com/agl/ed25519 v0.0.0-20170116200512-5312a6153412 h1:w1UutsfOrms1J05zt7I
|
||||
github.com/agl/ed25519 v0.0.0-20170116200512-5312a6153412/go.mod h1:WPjqKcmVOxf0XSf3YxCJs6N6AOSrOx3obionmG7T0y0=
|
||||
github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
|
||||
github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
|
||||
github.com/aliyun/aliyun-oss-go-sdk v2.0.0/go.mod h1:T/Aws4fEfogEE9v+HPhhw+CntffsBHJ8nXQCwKr0/g8=
|
||||
github.com/aliyun/aliyun-oss-go-sdk v2.0.5+incompatible h1:A3oZlWPD/Poa19FvNbw+Zu4yKAurDBTjlRDilYGBiS4=
|
||||
github.com/aliyun/aliyun-oss-go-sdk v2.0.5+incompatible/go.mod h1:T/Aws4fEfogEE9v+HPhhw+CntffsBHJ8nXQCwKr0/g8=
|
||||
github.com/apache/thrift v0.12.0/go.mod h1:cp2SuWMxlEZw2r+iP2GNCdIi4C1qmUzdZFSVb+bacwQ=
|
||||
@@ -248,6 +249,8 @@ golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8U
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/image v0.0.0-20190501045829-6d32002ffd75 h1:TbGuee8sSq15Iguxu4deQ7+Bqq/d2rsQejGcEtADAMQ=
|
||||
golang.org/x/image v0.0.0-20190501045829-6d32002ffd75/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
|
||||
golang.org/x/image v0.0.0-20211028202545-6944b10bf410 h1:hTftEOvwiOq2+O8k2D5/Q7COC7k5Qcrgc2TFURJYnvQ=
|
||||
golang.org/x/image v0.0.0-20211028202545-6944b10bf410/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM=
|
||||
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
|
||||
golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
|
||||
golang.org/x/lint v0.0.0-20190301231843-5614ed5bae6f/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
|
||||
@@ -288,6 +291,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
|
||||
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
|
||||
golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
|
||||
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c h1:fqgJT0MGcGpPgpWU7VRdRjuArfcOvC4AoJmILihzhDg=
|
||||
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/tools v0.0.0-20180828015842-6cd1fcedba52/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
|
||||
@@ -37,6 +37,7 @@ func SignRequired(authInstance auth.Auth) gin.HandlerFunc {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
+14
-3
@@ -90,15 +90,27 @@ func TestSignRequired(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request, _ = http.NewRequest("GET", "/test", nil)
|
||||
SignRequiredFunc := SignRequired(auth.HMACAuth{SecretKey: []byte(util.RandStringRunes(256))})
|
||||
authInstance := auth.HMACAuth{SecretKey: []byte(util.RandStringRunes(256))}
|
||||
SignRequiredFunc := SignRequired(authInstance)
|
||||
|
||||
// 鉴权失败
|
||||
SignRequiredFunc(c)
|
||||
asserts.NotNil(c)
|
||||
asserts.True(c.IsAborted())
|
||||
|
||||
c, _ = gin.CreateTestContext(rec)
|
||||
c.Request, _ = http.NewRequest("PUT", "/test", nil)
|
||||
SignRequiredFunc(c)
|
||||
asserts.NotNil(c)
|
||||
asserts.True(c.IsAborted())
|
||||
|
||||
// Sign verify success
|
||||
c, _ = gin.CreateTestContext(rec)
|
||||
c.Request, _ = http.NewRequest("PUT", "/test", nil)
|
||||
c.Request = auth.SignRequest(authInstance, c.Request, 0)
|
||||
SignRequiredFunc(c)
|
||||
asserts.NotNil(c)
|
||||
asserts.False(c.IsAborted())
|
||||
}
|
||||
|
||||
func TestWebDAVAuth(t *testing.T) {
|
||||
@@ -780,8 +792,6 @@ func TestS3CallbackAuth(t *testing.T) {
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "group_id"}).AddRow(1, 1))
|
||||
mock.ExpectQuery("SELECT(.+)groups(.+)").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "policies"}).AddRow(1, "[702]"))
|
||||
mock.ExpectQuery("SELECT(.+)policies(.+)").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "access_key", "secret_key"}).AddRow(2, "123", "123"))
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Params = []gin.Param{
|
||||
{"key", "testCallBackUpyun"},
|
||||
@@ -789,5 +799,6 @@ func TestS3CallbackAuth(t *testing.T) {
|
||||
c.Request, _ = http.NewRequest("POST", "/api/v3/callback/upyun/testCallBackUpyun", ioutil.NopCloser(strings.NewReader("1")))
|
||||
AuthFunc(c)
|
||||
asserts.False(c.IsAborted())
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ package middleware
|
||||
import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/slave"
|
||||
"github.com/gin-gonic/gin"
|
||||
"strconv"
|
||||
)
|
||||
@@ -11,19 +10,19 @@ import (
|
||||
// MasterMetadata 解析主机节点发来请求的包含主机节点信息的元数据
|
||||
func MasterMetadata() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Set("MasterSiteID", c.GetHeader("X-Site-Id"))
|
||||
c.Set("MasterSiteURL", c.GetHeader("X-Site-Url"))
|
||||
c.Set("MasterVersion", c.GetHeader("X-Cloudreve-Version"))
|
||||
c.Set("MasterSiteID", c.GetHeader("X-Cr-Site-Id"))
|
||||
c.Set("MasterSiteURL", c.GetHeader("X-Cr-Site-Url"))
|
||||
c.Set("MasterVersion", c.GetHeader("X-Cr-Cloudreve-Version"))
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// UseSlaveAria2Instance 从机用于获取对应主机节点的Aria2实例
|
||||
func UseSlaveAria2Instance() gin.HandlerFunc {
|
||||
func UseSlaveAria2Instance(clusterController cluster.Controller) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if siteID, exist := c.Get("MasterSiteID"); exist {
|
||||
// 获取对应主机节点的从机Aria2实例
|
||||
caller, err := slave.DefaultController.GetAria2Instance(siteID.(string))
|
||||
caller, err := clusterController.GetAria2Instance(siteID.(string))
|
||||
if err != nil {
|
||||
c.JSON(200, serializer.Err(serializer.CodeNotSet, "无法获取 Aria2 实例", err))
|
||||
c.Abort()
|
||||
@@ -40,16 +39,16 @@ func UseSlaveAria2Instance() gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
func SlaveRPCSignRequired() gin.HandlerFunc {
|
||||
func SlaveRPCSignRequired(nodePool cluster.Pool) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
nodeID, err := strconv.ParseUint(c.GetHeader("X-Node-Id"), 10, 64)
|
||||
nodeID, err := strconv.ParseUint(c.GetHeader("X-Cr-Node-Id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(200, serializer.ParamErr("未知的主机节点ID", err))
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
slaveNode := cluster.Default.GetNodeByID(uint(nodeID))
|
||||
slaveNode := nodePool.GetNodeByID(uint(nodeID))
|
||||
if slaveNode == nil {
|
||||
c.JSON(200, serializer.ParamErr("未知的主机节点ID", err))
|
||||
c.Abort()
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mocks/controllermock"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMasterMetadata(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
masterMetaDataFunc := MasterMetadata()
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
|
||||
c.Request.Header = map[string][]string{
|
||||
"X-Cr-Site-Id": {"expectedSiteID"},
|
||||
"X-Cr-Site-Url": {"expectedSiteURL"},
|
||||
"X-Cr-Cloudreve-Version": {"expectedMasterVersion"},
|
||||
}
|
||||
masterMetaDataFunc(c)
|
||||
siteID, _ := c.Get("MasterSiteID")
|
||||
siteURL, _ := c.Get("MasterSiteURL")
|
||||
siteVersion, _ := c.Get("MasterVersion")
|
||||
|
||||
a.Equal("expectedSiteID", siteID.(string))
|
||||
a.Equal("expectedSiteURL", siteURL.(string))
|
||||
a.Equal("expectedMasterVersion", siteVersion.(string))
|
||||
}
|
||||
|
||||
func TestSlaveRPCSignRequired(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
np := &cluster.NodePool{}
|
||||
np.Init()
|
||||
slaveRPCSignRequiredFunc := SlaveRPCSignRequired(np)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
// id parse failed
|
||||
{
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.Header.Set("X-Cr-Node-Id", "unknown")
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// node id not exist
|
||||
{
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.Header.Set("X-Cr-Node-Id", "38")
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
authInstance := auth.HMACAuth{SecretKey: []byte("")}
|
||||
np.Add(&model.Node{Model: gorm.Model{
|
||||
ID: 38,
|
||||
}})
|
||||
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest("POST", "/", nil)
|
||||
c.Request.Header.Set("X-Cr-Node-Id", "38")
|
||||
c.Request = auth.SignRequest(authInstance, c.Request, 0)
|
||||
slaveRPCSignRequiredFunc(c)
|
||||
a.False(c.IsAborted())
|
||||
}
|
||||
}
|
||||
|
||||
func TestUseSlaveAria2Instance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
|
||||
// MasterSiteID not set
|
||||
{
|
||||
testController := &controllermock.SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
}
|
||||
|
||||
// Cannot get aria2 instances
|
||||
{
|
||||
testController := &controllermock.SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Set("MasterSiteID", "expectedSiteID")
|
||||
testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, errors.New("error"))
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.True(c.IsAborted())
|
||||
testController.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Success
|
||||
{
|
||||
testController := &controllermock.SlaveControllerMock{}
|
||||
useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Set("MasterSiteID", "expectedSiteID")
|
||||
testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, nil)
|
||||
useSlaveAria2InstanceFunc(c)
|
||||
a.False(c.IsAborted())
|
||||
res, _ := c.Get("MasterAria2Instance")
|
||||
a.NotNil(res)
|
||||
testController.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
@@ -177,3 +177,14 @@ func TestDownload_Delete(t *testing.T) {
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestDownload_GetNodeID(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
record := Download{}
|
||||
|
||||
// compatible with 3.4
|
||||
a.EqualValues(1, record.GetNodeID())
|
||||
|
||||
record.NodeID = 5
|
||||
a.EqualValues(5, record.GetNodeID())
|
||||
}
|
||||
|
||||
@@ -1,12 +1,17 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"github.com/fatih/color"
|
||||
"github.com/gofrs/uuid"
|
||||
"github.com/hashicorp/go-version"
|
||||
"github.com/jinzhu/gorm"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// 是否需要迁移
|
||||
@@ -54,6 +59,9 @@ func migration() {
|
||||
// 向设置数据表添加初始设置
|
||||
addDefaultSettings()
|
||||
|
||||
// 执行数据库升级脚本
|
||||
execUpgradeScripts()
|
||||
|
||||
util.Log().Info("数据库初始化结束")
|
||||
|
||||
}
|
||||
@@ -181,6 +189,7 @@ Neue',Helvetica,Arial,sans-serif; box-sizing: border-box; font-size: 14px; verti
|
||||
{Name: "pwa_display", Value: "standalone", Type: "pwa"},
|
||||
{Name: "pwa_theme_color", Value: "#000000", Type: "pwa"},
|
||||
{Name: "pwa_background_color", Value: "#ffffff", Type: "pwa"},
|
||||
{Name: "office_preview_service", Value: "https://view.officeapps.live.com/op/view.aspx?src={$src}", Type: "preview"},
|
||||
}
|
||||
|
||||
for _, value := range defaultSettings {
|
||||
@@ -290,3 +299,17 @@ func addDefaultNode() {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func execUpgradeScripts() {
|
||||
s := invoker.ListPrefix("UpgradeTo")
|
||||
versions := make([]*version.Version, len(s))
|
||||
for i, raw := range s {
|
||||
v, _ := version.NewVersion(strings.TrimPrefix(raw, "UpgradeTo"))
|
||||
versions[i] = v
|
||||
}
|
||||
sort.Sort(version.Collection(versions))
|
||||
|
||||
for i := 0; i < len(versions); i++ {
|
||||
invoker.RunDBScript("UpgradeTo"+versions[i].String(), context.Background())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetNodeByID(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
res, err := GetNodeByID(1)
|
||||
a.NoError(err)
|
||||
a.EqualValues(1, res.ID)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestGetNodesByStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow(NodeActive))
|
||||
res, err := GetNodesByStatus(NodeActive)
|
||||
a.NoError(err)
|
||||
a.Len(res, 1)
|
||||
a.EqualValues(NodeActive, res[0].Status)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestNode_AfterFind(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
// No aria2 options
|
||||
{
|
||||
a.NoError(node.AfterFind())
|
||||
}
|
||||
|
||||
// with aria2 options
|
||||
{
|
||||
node.Aria2Options = `{"timeout":1}`
|
||||
a.NoError(node.AfterFind())
|
||||
a.Equal(1, node.Aria2OptionsSerialized.Timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNode_BeforeSave(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
node.Aria2OptionsSerialized.Timeout = 1
|
||||
a.NoError(node.BeforeSave())
|
||||
a.Contains(node.Aria2Options, "1")
|
||||
}
|
||||
|
||||
func TestNode_SetStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
node := &Node{}
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)nodes").WithArgs(NodeActive, sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
a.NoError(node.SetStatus(NodeActive))
|
||||
a.Equal(NodeActive, node.Status)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
@@ -327,3 +327,18 @@ func TestPolicy_IsThumbExist(t *testing.T) {
|
||||
asserts.Equal(testCase.expect, policy.IsThumbExist(testCase.name))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPolicy_UpdateAccessKeyAndClearCache(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
cache.Set("policy_1331", Policy{}, 3600)
|
||||
p := &Policy{}
|
||||
p.ID = 1331
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WithArgs("ak", sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.NoError(p.UpdateAccessKeyAndClearCache("ak"))
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
_, ok := cache.Get("policy_1331")
|
||||
a.False(ok)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
package scripts
|
||||
|
||||
import "github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
|
||||
|
||||
func Init() {
|
||||
invoker.Register("ResetAdminPassword", ResetAdminPassword(0))
|
||||
invoker.Register("CalibrateUserStorage", UserStorageCalibration(0))
|
||||
invoker.Register("UpgradeTo3.4.0", UpgradeTo340(0))
|
||||
}
|
||||
@@ -1,25 +0,0 @@
|
||||
package scripts
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
type DBScript interface {
|
||||
Run(ctx context.Context)
|
||||
}
|
||||
|
||||
var availableScripts = make(map[string]DBScript)
|
||||
|
||||
func RunDBScript(name string, ctx context.Context) error {
|
||||
if script, ok := availableScripts[name]; ok {
|
||||
script.Run(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("数据库脚本 [%s] 不存在", name)
|
||||
}
|
||||
|
||||
func register(name string, script DBScript) {
|
||||
availableScripts[name] = script
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package invoker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type DBScript interface {
|
||||
Run(ctx context.Context)
|
||||
}
|
||||
|
||||
var availableScripts = make(map[string]DBScript)
|
||||
|
||||
func RunDBScript(name string, ctx context.Context) error {
|
||||
if script, ok := availableScripts[name]; ok {
|
||||
util.Log().Info("开始执行数据库脚本 [%s]", name)
|
||||
script.Run(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("数据库脚本 [%s] 不存在", name)
|
||||
}
|
||||
|
||||
func Register(name string, script DBScript) {
|
||||
availableScripts[name] = script
|
||||
}
|
||||
|
||||
func ListPrefix(prefix string) []string {
|
||||
var scripts []string
|
||||
for name := range availableScripts {
|
||||
if strings.HasPrefix(name, prefix) {
|
||||
scripts = append(scripts, name)
|
||||
}
|
||||
}
|
||||
return scripts
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package invoker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type TestScript int
|
||||
|
||||
func (script TestScript) Run(ctx context.Context) {
|
||||
|
||||
}
|
||||
|
||||
func TestRunDBScript(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
Register("test", TestScript(0))
|
||||
|
||||
// 不存在
|
||||
{
|
||||
asserts.Error(RunDBScript("else", context.Background()))
|
||||
}
|
||||
|
||||
// 存在
|
||||
{
|
||||
asserts.NoError(RunDBScript("test", context.Background()))
|
||||
}
|
||||
}
|
||||
|
||||
func TestListPrefix(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
Register("U1", TestScript(0))
|
||||
Register("U2", TestScript(0))
|
||||
Register("U3", TestScript(0))
|
||||
Register("P1", TestScript(0))
|
||||
|
||||
res := ListPrefix("U")
|
||||
asserts.Len(res, 3)
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
package scripts
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
var mock sqlmock.Sqlmock
|
||||
var mockDB *gorm.DB
|
||||
|
||||
type TestScript int
|
||||
|
||||
func (script TestScript) Run(ctx context.Context) {
|
||||
|
||||
}
|
||||
|
||||
// TestMain 初始化数据库Mock
|
||||
func TestMain(m *testing.M) {
|
||||
var db *sql.DB
|
||||
var err error
|
||||
db, mock, err = sqlmock.New()
|
||||
if err != nil {
|
||||
panic("An error was not expected when opening a stub database connection")
|
||||
}
|
||||
model.DB, _ = gorm.Open("mysql", db)
|
||||
mockDB = model.DB
|
||||
defer db.Close()
|
||||
m.Run()
|
||||
}
|
||||
|
||||
func TestRunDBScript(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
register("test", TestScript(0))
|
||||
|
||||
// 不存在
|
||||
{
|
||||
asserts.Error(RunDBScript("else", context.Background()))
|
||||
}
|
||||
|
||||
// 存在
|
||||
{
|
||||
asserts.NoError(RunDBScript("test", context.Background()))
|
||||
}
|
||||
}
|
||||
@@ -9,10 +9,6 @@ import (
|
||||
|
||||
type ResetAdminPassword int
|
||||
|
||||
func init() {
|
||||
register("ResetAdminPassword", ResetAdminPassword(0))
|
||||
}
|
||||
|
||||
// Run 运行脚本从社区版升级至 Pro 版
|
||||
func (script ResetAdminPassword) Run(ctx context.Context) {
|
||||
// 查找用户
|
||||
|
||||
@@ -8,10 +8,6 @@ import (
|
||||
|
||||
type UserStorageCalibration int
|
||||
|
||||
func init() {
|
||||
register("CalibrateUserStorage", UserStorageCalibration(0))
|
||||
}
|
||||
|
||||
type storageResult struct {
|
||||
Total uint64
|
||||
}
|
||||
|
||||
@@ -2,11 +2,31 @@ package scripts
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
var mock sqlmock.Sqlmock
|
||||
var mockDB *gorm.DB
|
||||
|
||||
// TestMain 初始化数据库Mock
|
||||
func TestMain(m *testing.M) {
|
||||
var db *sql.DB
|
||||
var err error
|
||||
db, mock, err = sqlmock.New()
|
||||
if err != nil {
|
||||
panic("An error was not expected when opening a stub database connection")
|
||||
}
|
||||
model.DB, _ = gorm.Open("mysql", db)
|
||||
mockDB = model.DB
|
||||
defer db.Close()
|
||||
m.Run()
|
||||
}
|
||||
|
||||
func TestUserStorageCalibration_Run(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
script := UserStorageCalibration(0)
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package scripts
|
||||
|
||||
import (
|
||||
"context"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type UpgradeTo340 int
|
||||
|
||||
// Run upgrade from older version to 3.4.0
|
||||
func (script UpgradeTo340) Run(ctx context.Context) {
|
||||
// 取回老版本 aria2 设定
|
||||
old := model.GetSettingByType([]string{"aria2"})
|
||||
if len(old) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// 写入到新版本的节点设定
|
||||
n, err := model.GetNodeByID(1)
|
||||
if err != nil {
|
||||
util.Log().Error("找不到主机节点, %s", err)
|
||||
}
|
||||
|
||||
n.Aria2Enabled = old["aria2_rpcurl"] != ""
|
||||
n.Aria2OptionsSerialized.Options = old["aria2_options"]
|
||||
n.Aria2OptionsSerialized.Server = old["aria2_rpcurl"]
|
||||
|
||||
interval, err := strconv.Atoi(old["aria2_interval"])
|
||||
if err != nil {
|
||||
interval = 10
|
||||
}
|
||||
n.Aria2OptionsSerialized.Interval = interval
|
||||
n.Aria2OptionsSerialized.TempPath = old["aria2_temp_path"]
|
||||
n.Aria2OptionsSerialized.Token = old["aria2_token"]
|
||||
if err := model.DB.Save(&n).Error; err != nil {
|
||||
util.Log().Error("无法保存主机节点 Aria2 配置信息, %s", err)
|
||||
} else {
|
||||
model.DB.Where("type = ?", "aria2").Delete(model.Setting{})
|
||||
util.Log().Info("Aria2 配置信息已成功迁移至 3.4.0+ 版本的模式")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package scripts
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestUpgradeTo340_Run(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
script := UpgradeTo340(0)
|
||||
|
||||
// skip
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name"}))
|
||||
script.Run(context.Background())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// node not found
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name"}).AddRow("1"))
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
script.Run(context.Background())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name", "value"}).
|
||||
AddRow("aria2_rpcurl", "expected_aria2_rpcurl").
|
||||
AddRow("aria2_interval", "expected_aria2_interval").
|
||||
AddRow("aria2_temp_path", "expected_aria2_temp_path").
|
||||
AddRow("aria2_token", "expected_aria2_token").
|
||||
AddRow("aria2_options", "{}"))
|
||||
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
script.Run(context.Background())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// failed
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name", "value"}).
|
||||
AddRow("aria2_rpcurl", "expected_aria2_rpcurl").
|
||||
AddRow("aria2_interval", "expected_aria2_interval").
|
||||
AddRow("aria2_temp_path", "expected_aria2_temp_path").
|
||||
AddRow("aria2_token", "expected_aria2_token").
|
||||
AddRow("aria2_options", "{}"))
|
||||
|
||||
mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
mock.ExpectRollback()
|
||||
script.Run(context.Background())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
}
|
||||
@@ -188,18 +188,6 @@ func TestShare_CanBeDownloadBy(t *testing.T) {
|
||||
asserts.Error(share.CanBeDownloadBy(user))
|
||||
}
|
||||
|
||||
// 未登录,需要积分
|
||||
{
|
||||
user := &User{
|
||||
Group: Group{
|
||||
OptionsSerialized: GroupOption{
|
||||
ShareDownload: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
asserts.Error(share.CanBeDownloadBy(user))
|
||||
}
|
||||
|
||||
// 成功
|
||||
{
|
||||
user := &User{
|
||||
|
||||
+3
-3
@@ -55,9 +55,9 @@ func TestGetTagsByUID(t *testing.T) {
|
||||
|
||||
func TestGetTagsByID(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
res, err := GetTasksByID(1)
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"name"}).AddRow("tag"))
|
||||
res, err := GetTagsByUID(1)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.NoError(err)
|
||||
asserts.EqualValues(1, res.ID)
|
||||
asserts.EqualValues("tag", res[0].Name)
|
||||
}
|
||||
|
||||
@@ -91,3 +91,14 @@ func TestListTasks(t *testing.T) {
|
||||
asserts.EqualValues(5, total)
|
||||
asserts.Len(res, 1)
|
||||
}
|
||||
|
||||
func TestGetTasksByStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
|
||||
mock.ExpectQuery("SELECT(.+)").
|
||||
WithArgs(1, 2).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
res := GetTasksByStatus(1, 2)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.Len(res, 1)
|
||||
}
|
||||
|
||||
+1
-1
@@ -35,7 +35,7 @@ type User struct {
|
||||
Storage uint64
|
||||
TwoFactor string
|
||||
Avatar string
|
||||
Options string `json:"-",gorm:"type:text"`
|
||||
Options string `json:"-" gorm:"type:text"`
|
||||
Authn string `gorm:"type:text"`
|
||||
|
||||
// 关联模型
|
||||
|
||||
+2
-11
@@ -177,10 +177,10 @@ func TestNewUser(t *testing.T) {
|
||||
|
||||
func TestUser_AfterFind(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
cache.Deletes([]string{"1"}, "policy_")
|
||||
cache.Deletes([]string{"0"}, "policy_")
|
||||
|
||||
policyRows := sqlmock.NewRows([]string{"id", "name"}).
|
||||
AddRow(1, "默认存储策略")
|
||||
AddRow(144, "默认存储策略")
|
||||
mock.ExpectQuery("^SELECT (.+)").WillReturnRows(policyRows)
|
||||
|
||||
newUser := NewUser()
|
||||
@@ -240,11 +240,6 @@ func TestUser_GetRemainingCapacity(t *testing.T) {
|
||||
newUser.Group.MaxStorage = 100
|
||||
newUser.Storage = 200
|
||||
asserts.Equal(uint64(0), newUser.GetRemainingCapacity())
|
||||
|
||||
cache.Set("pack_size_0", uint64(10), 0)
|
||||
newUser.Group.MaxStorage = 100
|
||||
newUser.Storage = 101
|
||||
asserts.Equal(uint64(9), newUser.GetRemainingCapacity())
|
||||
}
|
||||
|
||||
func TestUser_DeductionCapacity(t *testing.T) {
|
||||
@@ -280,10 +275,6 @@ func TestUser_DeductionCapacity(t *testing.T) {
|
||||
asserts.Equal(false, newUser.IncreaseStorage(1))
|
||||
asserts.Equal(uint64(100), newUser.Storage)
|
||||
|
||||
cache.Set("pack_size_1", uint64(1), 0)
|
||||
asserts.Equal(true, newUser.IncreaseStorage(1))
|
||||
asserts.Equal(uint64(101), newUser.Storage)
|
||||
|
||||
asserts.True(newUser.IncreaseStorage(0))
|
||||
}
|
||||
|
||||
|
||||
@@ -55,6 +55,6 @@ func TestDeleteWebDAVAccountByID(t *testing.T) {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
asserts.NoError(DeleteTagByID(1, 1))
|
||||
DeleteWebDAVAccountByID(1, 1)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
+4
-2
@@ -3,6 +3,8 @@ package aria2
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"net/url"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -31,7 +33,7 @@ func GetLoadBalancer() balancer.Balancer {
|
||||
}
|
||||
|
||||
// Init 初始化
|
||||
func Init(isReload bool) {
|
||||
func Init(isReload bool, pool cluster.Pool, mqClient mq.MQ) {
|
||||
Lock.Lock()
|
||||
LB = balancer.NewBalancer("RoundRobin")
|
||||
Lock.Unlock()
|
||||
@@ -42,7 +44,7 @@ func Init(isReload bool) {
|
||||
|
||||
for i := 0; i < len(unfinished); i++ {
|
||||
// 创建任务监控
|
||||
monitor.NewMonitor(&unfinished[i])
|
||||
monitor.NewMonitor(&unfinished[i], pool, mqClient)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+30
-56
@@ -2,14 +2,15 @@ package aria2
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mocks"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/monitor"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
var mock sqlmock.Sqlmock
|
||||
@@ -27,66 +28,39 @@ func TestMain(m *testing.M) {
|
||||
m.Run()
|
||||
}
|
||||
|
||||
func TestDummyAria2(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
instance := DummyAria2{}
|
||||
asserts.Error(instance.CreateTask(nil, nil))
|
||||
_, err := instance.Status(nil)
|
||||
asserts.Error(err)
|
||||
asserts.Error(instance.Cancel(nil))
|
||||
asserts.Error(instance.Select(nil, nil))
|
||||
}
|
||||
|
||||
func TestInit(t *testing.T) {
|
||||
monitor.MAX_RETRY = 0
|
||||
asserts := assert.New(t)
|
||||
cache.Set("setting_aria2_token", "1", 0)
|
||||
cache.Set("setting_aria2_call_timeout", "5", 0)
|
||||
cache.Set("setting_aria2_options", `[]`, 0)
|
||||
a := assert.New(t)
|
||||
mockPool := &mocks.NodePoolMock{}
|
||||
mockPool.On("GetNodeByID", testMock.Anything).Return(nil)
|
||||
mockQueue := mq.NewMQ()
|
||||
|
||||
// 未指定RPC地址,跳过
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
Init(false, mockPool, mockQueue)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockPool.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestTestRPCConnection(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
|
||||
// url not legal
|
||||
{
|
||||
cache.Set("setting_aria2_rpcurl", "", 0)
|
||||
Init(false)
|
||||
asserts.IsType(&DummyAria2{}, Instance)
|
||||
res, err := TestRPCConnection(string([]byte{0x7f}), "", 10)
|
||||
a.Error(err)
|
||||
a.Empty(res.Version)
|
||||
}
|
||||
|
||||
// 无法解析服务器地址
|
||||
// rpc failed
|
||||
{
|
||||
cache.Set("setting_aria2_rpcurl", string(byte(0x7f)), 0)
|
||||
Init(false)
|
||||
asserts.IsType(&DummyAria2{}, Instance)
|
||||
}
|
||||
|
||||
// 无法解析全局配置
|
||||
{
|
||||
Instance = &RPCService{}
|
||||
cache.Set("setting_aria2_options", "?", 0)
|
||||
cache.Set("setting_aria2_rpcurl", "ws://127.0.0.1:1234", 0)
|
||||
Init(false)
|
||||
asserts.IsType(&DummyAria2{}, Instance)
|
||||
}
|
||||
|
||||
// 连接失败
|
||||
{
|
||||
cache.Set("setting_aria2_options", "{}", 0)
|
||||
cache.Set("setting_aria2_rpcurl", "http://127.0.0.1:1234", 0)
|
||||
cache.Set("setting_aria2_call_timeout", "1", 0)
|
||||
cache.Set("setting_aria2_interval", "100", 0)
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"g_id"}).AddRow("1"))
|
||||
Init(false)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.IsType(&RPCService{}, Instance)
|
||||
res, err := TestRPCConnection("ws://0.0.0.0", "", 0)
|
||||
a.Error(err)
|
||||
a.Empty(res.Version)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetStatus(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
asserts.Equal(4, GetStatus("complete"))
|
||||
asserts.Equal(1, GetStatus("active"))
|
||||
asserts.Equal(0, GetStatus("waiting"))
|
||||
asserts.Equal(2, GetStatus("paused"))
|
||||
asserts.Equal(3, GetStatus("error"))
|
||||
asserts.Equal(5, GetStatus("removed"))
|
||||
asserts.Equal(6, GetStatus("?"))
|
||||
func TestGetLoadBalancer(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
a.NotPanics(func() {
|
||||
GetLoadBalancer()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
package aria2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
)
|
||||
|
||||
// RPCService 通过RPC服务的Aria2任务管理器
|
||||
type RPCService struct {
|
||||
options *clientOptions
|
||||
Caller rpc.Client
|
||||
}
|
||||
|
||||
type clientOptions struct {
|
||||
Options map[string]interface{} // 创建下载时额外添加的设置
|
||||
}
|
||||
|
||||
// Init 初始化
|
||||
func (client *RPCService) Init(server, secret string, timeout int, options map[string]interface{}) error {
|
||||
// 客户端已存在,则关闭先前连接
|
||||
if client.Caller != nil {
|
||||
client.Caller.Close()
|
||||
}
|
||||
|
||||
client.options = &clientOptions{
|
||||
Options: options,
|
||||
}
|
||||
caller, err := rpc.New(context.Background(), server, secret, time.Duration(timeout)*time.Second,
|
||||
mq.GlobalMQ)
|
||||
client.Caller = caller
|
||||
return err
|
||||
}
|
||||
|
||||
// Status 查询下载状态
|
||||
func (client *RPCService) Status(task *model.Download) (rpc.StatusInfo, error) {
|
||||
res, err := client.Caller.TellStatus(task.GID)
|
||||
if err != nil {
|
||||
// 失败后重试
|
||||
util.Log().Debug("无法获取离线下载状态,%s,10秒钟后重试", err)
|
||||
time.Sleep(time.Duration(10) * time.Second)
|
||||
res, err = client.Caller.TellStatus(task.GID)
|
||||
}
|
||||
|
||||
return res, err
|
||||
}
|
||||
|
||||
// Cancel 取消下载
|
||||
func (client *RPCService) Cancel(task *model.Download) error {
|
||||
// 取消下载任务
|
||||
_, err := client.Caller.Remove(task.GID)
|
||||
if err != nil {
|
||||
util.Log().Warning("无法取消离线下载任务[%s], %s", task.GID, err)
|
||||
}
|
||||
|
||||
//// 删除临时文件
|
||||
//util.Log().Debug("离线下载任务[%s]已取消,1 分钟后删除临时文件", task.GID)
|
||||
//go func(task *model.Download) {
|
||||
// select {
|
||||
// case <-time.After(time.Duration(60) * time.Second):
|
||||
// err := os.RemoveAll(task.Parent)
|
||||
// if err != nil {
|
||||
// util.Log().Warning("无法删除离线下载临时目录[%s], %s", task.Parent, err)
|
||||
// }
|
||||
// }
|
||||
//}(task)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// Select 选取要下载的文件
|
||||
func (client *RPCService) Select(task *model.Download, files []int) error {
|
||||
var selected = make([]string, len(files))
|
||||
for i := 0; i < len(files); i++ {
|
||||
selected[i] = strconv.Itoa(files[i])
|
||||
}
|
||||
_, err := client.Caller.ChangeOption(task.GID, map[string]interface{}{"select-file": strings.Join(selected, ",")})
|
||||
return err
|
||||
}
|
||||
|
||||
// CreateTask 创建新任务
|
||||
func (client *RPCService) CreateTask(task *model.Download, groupOptions map[string]interface{}) (string, error) {
|
||||
// 生成存储路径
|
||||
path := filepath.Join(
|
||||
model.GetSettingByName("aria2_temp_path"),
|
||||
"aria2",
|
||||
strconv.FormatInt(time.Now().UnixNano(), 10),
|
||||
)
|
||||
|
||||
// 创建下载任务
|
||||
options := map[string]interface{}{
|
||||
"dir": path,
|
||||
}
|
||||
for k, v := range client.options.Options {
|
||||
options[k] = v
|
||||
}
|
||||
for k, v := range groupOptions {
|
||||
options[k] = v
|
||||
}
|
||||
|
||||
gid, err := client.Caller.AddURI(task.Source, options)
|
||||
if err != nil || gid == "" {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return gid, nil
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
package aria2
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestRPCService_Init(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
caller := &RPCService{}
|
||||
asserts.Error(caller.Init("ws://", "", 1, nil))
|
||||
asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
|
||||
}
|
||||
|
||||
func TestRPCService_Status(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
caller := &RPCService{}
|
||||
asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
|
||||
|
||||
_, err := caller.Status(&model.Download{})
|
||||
asserts.Error(err)
|
||||
}
|
||||
|
||||
func TestRPCService_Cancel(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
caller := &RPCService{}
|
||||
asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
|
||||
|
||||
err := caller.Cancel(&model.Download{Parent: "test"})
|
||||
asserts.Error(err)
|
||||
}
|
||||
|
||||
func TestRPCService_Select(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
caller := &RPCService{}
|
||||
asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
|
||||
|
||||
err := caller.Select(&model.Download{Parent: "test"}, []int{1, 2, 3})
|
||||
asserts.Error(err)
|
||||
}
|
||||
|
||||
func TestRPCService_CreateTask(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
caller := &RPCService{}
|
||||
asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
|
||||
cache.Set("setting_aria2_temp_path", "test", 0)
|
||||
err := caller.CreateTask(&model.Download{Parent: "test"}, map[string]interface{}{"1": "1"})
|
||||
asserts.Error(err)
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDummyAria2(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
d := &DummyAria2{}
|
||||
|
||||
a.NoError(d.Init())
|
||||
|
||||
res, err := d.CreateTask(&model.Download{}, map[string]interface{}{})
|
||||
a.Empty(res)
|
||||
a.Error(err)
|
||||
|
||||
_, err = d.Status(&model.Download{})
|
||||
a.Error(err)
|
||||
|
||||
err = d.Cancel(&model.Download{})
|
||||
a.Error(err)
|
||||
|
||||
err = d.Select(&model.Download{}, []int{})
|
||||
a.Error(err)
|
||||
|
||||
configRes := d.GetConfig()
|
||||
a.NotNil(configRes)
|
||||
|
||||
err = d.DeleteTempFile(&model.Download{})
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
func TestGetStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
|
||||
a.Equal(GetStatus("complete"), Complete)
|
||||
a.Equal(GetStatus("active"), Downloading)
|
||||
a.Equal(GetStatus("waiting"), Ready)
|
||||
a.Equal(GetStatus("paused"), Paused)
|
||||
a.Equal(GetStatus("error"), Error)
|
||||
a.Equal(GetStatus("removed"), Canceled)
|
||||
a.Equal(GetStatus("unknown"), Unknown)
|
||||
}
|
||||
@@ -33,29 +33,29 @@ type Monitor struct {
|
||||
var MAX_RETRY = 10
|
||||
|
||||
// NewMonitor 新建离线下载状态监控
|
||||
func NewMonitor(task *model.Download) {
|
||||
func NewMonitor(task *model.Download, pool cluster.Pool, mqClient mq.MQ) {
|
||||
monitor := &Monitor{
|
||||
Task: task,
|
||||
notifier: make(chan mq.Message),
|
||||
node: cluster.Default.GetNodeByID(task.GetNodeID()),
|
||||
node: pool.GetNodeByID(task.GetNodeID()),
|
||||
}
|
||||
|
||||
if monitor.node != nil {
|
||||
monitor.Interval = time.Duration(monitor.node.GetAria2Instance().GetConfig().Interval) * time.Second
|
||||
go monitor.Loop()
|
||||
go monitor.Loop(mqClient)
|
||||
|
||||
monitor.notifier = mq.GlobalMQ.Subscribe(monitor.Task.GID, 0)
|
||||
monitor.notifier = mqClient.Subscribe(monitor.Task.GID, 0)
|
||||
} else {
|
||||
monitor.setErrorStatus(errors.New("节点不可用"))
|
||||
}
|
||||
}
|
||||
|
||||
// Loop 开启监控循环
|
||||
func (monitor *Monitor) Loop() {
|
||||
defer mq.GlobalMQ.Unsubscribe(monitor.Task.GID, monitor.notifier)
|
||||
func (monitor *Monitor) Loop(mqClient mq.MQ) {
|
||||
defer mqClient.Unsubscribe(monitor.Task.GID, monitor.notifier)
|
||||
|
||||
// 首次循环立即更新
|
||||
interval := time.Duration(0)
|
||||
interval := 50 * time.Millisecond
|
||||
|
||||
for {
|
||||
select {
|
||||
@@ -112,7 +112,7 @@ func (monitor *Monitor) Update() bool {
|
||||
|
||||
switch status.Status {
|
||||
case "complete":
|
||||
return monitor.Complete(status)
|
||||
return monitor.Complete(task.TaskPoll)
|
||||
case "error":
|
||||
return monitor.Error(status)
|
||||
case "active", "waiting", "paused":
|
||||
@@ -235,7 +235,7 @@ func (monitor *Monitor) RemoveTempFolder() {
|
||||
}
|
||||
|
||||
// Complete 完成下载,返回是否中断监控
|
||||
func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
|
||||
func (monitor *Monitor) Complete(pool task.Pool) bool {
|
||||
// 创建中转任务
|
||||
file := make([]string, 0, len(monitor.Task.StatusInfo.Files))
|
||||
sizes := make(map[string]uint64, len(monitor.Task.StatusInfo.Files))
|
||||
@@ -259,11 +259,12 @@ func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
|
||||
)
|
||||
if err != nil {
|
||||
monitor.setErrorStatus(err)
|
||||
monitor.RemoveTempFolder()
|
||||
return true
|
||||
}
|
||||
|
||||
// 提交中转任务
|
||||
task.TaskPoll.Submit(job)
|
||||
pool.Submit(job)
|
||||
|
||||
// 更新任务ID
|
||||
monitor.Task.TaskID = job.Model().ID
|
||||
|
||||
+366
-254
@@ -1,326 +1,438 @@
|
||||
package monitor
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mocks"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type InstanceMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
var mock sqlmock.Sqlmock
|
||||
|
||||
func (m InstanceMock) CreateTask(task *model.Download, options map[string]interface{}) error {
|
||||
args := m.Called(task, options)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Status(task *model.Download) (rpc.StatusInfo, error) {
|
||||
args := m.Called(task)
|
||||
return args.Get(0).(rpc.StatusInfo), args.Error(1)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Cancel(task *model.Download) error {
|
||||
args := m.Called(task)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m InstanceMock) Select(task *model.Download, files []int) error {
|
||||
args := m.Called(task, files)
|
||||
return args.Error(0)
|
||||
// TestMain 初始化数据库Mock
|
||||
func TestMain(m *testing.M) {
|
||||
var db *sql.DB
|
||||
var err error
|
||||
db, mock, err = sqlmock.New()
|
||||
if err != nil {
|
||||
panic("An error was not expected when opening a stub database connection")
|
||||
}
|
||||
model.DB, _ = gorm.Open("mysql", db)
|
||||
defer db.Close()
|
||||
m.Run()
|
||||
}
|
||||
|
||||
func TestNewMonitor(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
NewMonitor(&model.Download{GID: "gid"})
|
||||
_, ok := common.EventNotifier.Subscribes.Load("gid")
|
||||
asserts.True(ok)
|
||||
a := assert.New(t)
|
||||
mockMQ := mq.NewMQ()
|
||||
|
||||
// node not available
|
||||
{
|
||||
mockPool := &mocks.NodePoolMock{}
|
||||
mockPool.On("GetNodeByID", uint(1)).Return(nil)
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
task := &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
}
|
||||
NewMonitor(task, mockPool, mockMQ)
|
||||
mockPool.AssertExpectations(t)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.NotEmpty(task.Error)
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
|
||||
mockPool := &mocks.NodePoolMock{}
|
||||
mockPool.On("GetNodeByID", uint(1)).Return(mockNode)
|
||||
|
||||
task := &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
}
|
||||
NewMonitor(task, mockPool, mockMQ)
|
||||
mockNode.AssertExpectations(t)
|
||||
mockPool.AssertExpectations(t)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestMonitor_Loop(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
notifier := make(chan common.StatusEvent)
|
||||
MAX_RETRY = 0
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{GID: "gid"},
|
||||
Interval: time.Duration(1) * time.Second,
|
||||
notifier: notifier,
|
||||
}
|
||||
asserts.NotPanics(func() {
|
||||
monitor.Loop()
|
||||
})
|
||||
}
|
||||
|
||||
func TestMonitor_Update(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_Update",
|
||||
},
|
||||
Interval: time.Duration(1) * time.Second,
|
||||
a := assert.New(t)
|
||||
mockMQ := mq.NewMQ()
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
|
||||
m := &Monitor{
|
||||
retried: MAX_RETRY,
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
notifier: mockMQ.Subscribe("test", 1),
|
||||
}
|
||||
|
||||
// 无法获取状态
|
||||
// into interval loop
|
||||
{
|
||||
MAX_RETRY = 1
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, errors.New("error"))
|
||||
file, _ := util.CreatNestedFile("TestMonitor_Update/1")
|
||||
file.Close()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.True(monitor.Update())
|
||||
testInstance.AssertExpectations(t)
|
||||
asserts.False(util.Exists("TestMonitor_Update"))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
m.Loop(mockMQ)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
|
||||
// 磁力链下载重定向
|
||||
// into notifier loop
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{
|
||||
FollowedBy: []string{"1"},
|
||||
}, nil)
|
||||
monitor.Task.ID = 1
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
asserts.EqualValues("1", monitor.Task.GID)
|
||||
}
|
||||
|
||||
// 无法更新任务信息
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, nil)
|
||||
monitor.Task.ID = 1
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
aria2.mock.ExpectRollback()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回未知状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "?"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回被取消状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "removed"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回活跃状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "active"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.False(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回错误状态
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "error"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// 返回完成
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "complete"}, nil)
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
aria2.Instance = testInstance
|
||||
asserts.True(monitor.Update())
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
testInstance.AssertExpectations(t)
|
||||
m.Task.Error = ""
|
||||
mockMQ.Publish("test", mq.Message{})
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
m.Loop(mockMQ)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateTaskInfo(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_UpdateTaskInfo",
|
||||
},
|
||||
func TestMonitor_UpdateFailedAfterRetry(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
for i := 0; i < MAX_RETRY; i++ {
|
||||
a.False(m.Update())
|
||||
}
|
||||
|
||||
// 失败
|
||||
{
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
aria2.mock.ExpectRollback()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.Error(err)
|
||||
mockNode.AssertExpectations(t)
|
||||
a.True(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateMagentoFollow(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
FollowedBy: []string{"next"},
|
||||
}, nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.False(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
a.Equal("next", m.Task.GID)
|
||||
mockAria2.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateFailedToUpdateInfo(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{}, nil)
|
||||
mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
|
||||
mock.ExpectRollback()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateCompleted(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
Status: "complete",
|
||||
}, nil)
|
||||
mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
mockNode.On("ID").Return(uint(1))
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateError(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
Status: "error",
|
||||
ErrorMessage: "error",
|
||||
}, nil)
|
||||
mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
a.NotEmpty(m.Task.Error)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateActive(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
Status: "active",
|
||||
}, nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.False(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateRemoved(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
Status: "removed",
|
||||
}, nil)
|
||||
mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Update())
|
||||
a.Equal(common.Canceled, m.Task.Status)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateUnknown(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockAria2 := &mocks.Aria2Mock{}
|
||||
mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
|
||||
Status: "unknown",
|
||||
}, nil)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(mockAria2)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Update())
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockAria2.AssertExpectations(t)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestMonitor_UpdateTaskInfoValidateFailed(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
status := rpc.StatusInfo{
|
||||
Status: "completed",
|
||||
TotalLength: "100",
|
||||
CompletedLength: "50",
|
||||
DownloadSpeed: "20",
|
||||
}
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{Model: gorm.Model{ID: 1}},
|
||||
}
|
||||
|
||||
// 更新成功,无需校验
|
||||
{
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.NoError(err)
|
||||
}
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
// 更新成功,大小改变,需要校验,校验失败
|
||||
{
|
||||
testInstance := new(InstanceMock)
|
||||
testInstance.On("SlaveCancel", testMock.Anything).Return(nil)
|
||||
aria2.Instance = testInstance
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
err := monitor.UpdateTaskInfo(rpc.StatusInfo{TotalLength: "1"})
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
asserts.Error(err)
|
||||
testInstance.AssertExpectations(t)
|
||||
}
|
||||
err := m.UpdateTaskInfo(status)
|
||||
a.Error(err)
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestMonitor_ValidateFile(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
a := assert.New(t)
|
||||
m := &Monitor{
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_ValidateFile",
|
||||
Model: gorm.Model{ID: 1},
|
||||
TotalSize: 100,
|
||||
},
|
||||
}
|
||||
|
||||
// 无法创建文件系统
|
||||
// failed to create filesystem
|
||||
{
|
||||
monitor.Task.User = &model.User{
|
||||
m.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "unknown",
|
||||
Type: "random",
|
||||
},
|
||||
}
|
||||
asserts.Error(monitor.ValidateFile())
|
||||
a.Equal(filesystem.ErrUnknownPolicyType, m.ValidateFile())
|
||||
}
|
||||
|
||||
// 文件大小超出容量配额
|
||||
// User capacity not enough
|
||||
{
|
||||
cache.Set("pack_size_0", uint64(0), 0)
|
||||
monitor.Task.TotalSize = 11
|
||||
monitor.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "mock",
|
||||
},
|
||||
m.Task.User = &model.User{
|
||||
Group: model.Group{
|
||||
MaxStorage: 10,
|
||||
MaxStorage: 99,
|
||||
},
|
||||
Policy: model.Policy{
|
||||
Type: "local",
|
||||
},
|
||||
}
|
||||
asserts.Equal(filesystem.ErrInsufficientCapacity, monitor.ValidateFile())
|
||||
a.Equal(filesystem.ErrInsufficientCapacity, m.ValidateFile())
|
||||
}
|
||||
|
||||
// 单文件大小超出容量配额
|
||||
// single file too big
|
||||
{
|
||||
cache.Set("pack_size_0", uint64(0), 0)
|
||||
monitor.Task.TotalSize = 10
|
||||
monitor.Task.StatusInfo.Files = []rpc.FileInfo{
|
||||
m.Task.StatusInfo.Files = []rpc.FileInfo{
|
||||
{
|
||||
Length: "100",
|
||||
Selected: "true",
|
||||
Length: "6",
|
||||
},
|
||||
}
|
||||
monitor.Task.User = &model.User{
|
||||
Policy: model.Policy{
|
||||
Type: "mock",
|
||||
MaxSize: 5,
|
||||
},
|
||||
m.Task.User = &model.User{
|
||||
Group: model.Group{
|
||||
MaxStorage: 10,
|
||||
MaxStorage: 100,
|
||||
},
|
||||
Policy: model.Policy{
|
||||
Type: "local",
|
||||
MaxSize: 99,
|
||||
},
|
||||
}
|
||||
asserts.Equal(filesystem.ErrFileSizeTooBig, monitor.ValidateFile())
|
||||
a.Equal(filesystem.ErrFileSizeTooBig, m.ValidateFile())
|
||||
}
|
||||
|
||||
// all pass
|
||||
{
|
||||
m.Task.StatusInfo.Files = []rpc.FileInfo{
|
||||
{
|
||||
Length: "100",
|
||||
Selected: "true",
|
||||
},
|
||||
}
|
||||
m.Task.User = &model.User{
|
||||
Group: model.Group{
|
||||
MaxStorage: 100,
|
||||
},
|
||||
Policy: model.Policy{
|
||||
Type: "local",
|
||||
MaxSize: 100,
|
||||
},
|
||||
}
|
||||
a.NoError(m.ValidateFile())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitor_Complete(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
monitor := &Monitor{
|
||||
a := assert.New(t)
|
||||
mockNode := &mocks.NodeMock{}
|
||||
mockNode.On("ID").Return(uint(1))
|
||||
mockPool := &mocks.TaskPoolMock{}
|
||||
mockPool.On("Submit", testMock.Anything)
|
||||
m := &Monitor{
|
||||
node: mockNode,
|
||||
Task: &model.Download{
|
||||
Model: gorm.Model{ID: 1},
|
||||
GID: "gid",
|
||||
Parent: "TestMonitor_Complete",
|
||||
StatusInfo: rpc.StatusInfo{
|
||||
Files: []rpc.FileInfo{
|
||||
{
|
||||
Selected: "true",
|
||||
Path: "TestMonitor_Complete",
|
||||
},
|
||||
},
|
||||
},
|
||||
Model: gorm.Model{ID: 1},
|
||||
TotalSize: 100,
|
||||
UserID: 9414,
|
||||
},
|
||||
}
|
||||
m.Task.StatusInfo.Files = []rpc.FileInfo{
|
||||
{
|
||||
Length: "100",
|
||||
Selected: "true",
|
||||
},
|
||||
}
|
||||
|
||||
cache.Set("setting_max_worker_num", "1", 0)
|
||||
aria2.mock.ExpectQuery("SELECT(.+)tasks").WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
task.Init()
|
||||
aria2.mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
aria2.mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(9414))
|
||||
|
||||
aria2.mock.ExpectBegin()
|
||||
aria2.mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
aria2.mock.ExpectCommit()
|
||||
asserts.True(monitor.Complete(rpc.StatusInfo{}))
|
||||
asserts.NoError(aria2.mock.ExpectationsWereMet())
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
|
||||
a.True(m.Complete(mockPool))
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
mockNode.AssertExpectations(t)
|
||||
mockPool.AssertExpectations(t)
|
||||
}
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
package aria2
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestNotifier_Notify(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
notifier2 := &Notifier{}
|
||||
notifyChan := make(chan StatusEvent, 10)
|
||||
notifier2.Subscribe(notifyChan, "1")
|
||||
|
||||
// 未订阅
|
||||
{
|
||||
notifier2.Notify([]rpc.Event{rpc.Event{Gid: ""}}, 1)
|
||||
asserts.Len(notifyChan, 0)
|
||||
}
|
||||
|
||||
// 订阅
|
||||
{
|
||||
notifier2.Notify([]rpc.Event{{Gid: "1"}}, 1)
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnBtDownloadComplete([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnDownloadStart([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnDownloadPause([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnDownloadStop([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnDownloadComplete([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
|
||||
notifier2.OnDownloadError([]rpc.Event{{Gid: "1"}})
|
||||
asserts.Len(notifyChan, 1)
|
||||
<-notifyChan
|
||||
}
|
||||
}
|
||||
+8
-6
@@ -17,8 +17,10 @@ import (
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuthFailed = serializer.NewError(serializer.CodeNoPermissionErr, "鉴权失败", nil)
|
||||
ErrExpired = serializer.NewError(serializer.CodeSignExpired, "签名已过期", nil)
|
||||
ErrAuthFailed = serializer.NewError(serializer.CodeNoPermissionErr, "鉴权失败", nil)
|
||||
ErrAuthHeaderMissing = serializer.NewError(serializer.CodeNoPermissionErr, "authorization header is missing", nil)
|
||||
ErrExpiresMissing = serializer.NewError(serializer.CodeNoPermissionErr, "expire timestamp is missing", nil)
|
||||
ErrExpired = serializer.NewError(serializer.CodeSignExpired, "签名已过期", nil)
|
||||
)
|
||||
|
||||
// General 通用的认证接口
|
||||
@@ -33,7 +35,7 @@ type Auth interface {
|
||||
}
|
||||
|
||||
// SignRequest 对PUT\POST等复杂HTTP请求签名,只会对URI部分、
|
||||
// 请求正文、`X-`开头的header进行签名
|
||||
// 请求正文、`X-Cr-`开头的header进行签名
|
||||
func SignRequest(instance Auth, r *http.Request, expires int64) *http.Request {
|
||||
// 处理有效期
|
||||
if expires > 0 {
|
||||
@@ -55,7 +57,7 @@ func CheckRequest(instance Auth, r *http.Request) error {
|
||||
ok bool
|
||||
)
|
||||
if sign, ok = r.Header["Authorization"]; !ok || len(sign) == 0 {
|
||||
return ErrAuthFailed
|
||||
return ErrAuthHeaderMissing
|
||||
}
|
||||
sign[0] = strings.TrimPrefix(sign[0], "Bearer ")
|
||||
|
||||
@@ -67,7 +69,7 @@ func CheckRequest(instance Auth, r *http.Request) error {
|
||||
func getSignContent(r *http.Request) (rawSignString string) {
|
||||
// 读取所有body正文
|
||||
var body = []byte{}
|
||||
if _, ok := r.Header["X-Policy"]; !ok {
|
||||
if _, ok := r.Header["X-Cr-Policy"]; !ok {
|
||||
if r.Body != nil {
|
||||
body, _ = ioutil.ReadAll(r.Body)
|
||||
_ = r.Body.Close()
|
||||
@@ -78,7 +80,7 @@ func getSignContent(r *http.Request) (rawSignString string) {
|
||||
// 决定要签名的header
|
||||
var signedHeader []string
|
||||
for k, _ := range r.Header {
|
||||
if strings.HasPrefix(k, "X-") && k != "X-Filename" {
|
||||
if strings.HasPrefix(k, "X-Cr-") && k != "X-Cr-Filename" {
|
||||
signedHeader = append(signedHeader, fmt.Sprintf("%s=%s", k, r.Header.Get(k)))
|
||||
}
|
||||
}
|
||||
|
||||
+15
-2
@@ -70,7 +70,7 @@ func TestSignRequest(t *testing.T) {
|
||||
strings.NewReader("I am body."),
|
||||
)
|
||||
asserts.NoError(err)
|
||||
req.Header["X-Policy"] = []string{"I am Policy"}
|
||||
req.Header["X-Cr-Policy"] = []string{"I am Policy"}
|
||||
req = SignRequest(General, req, 10)
|
||||
asserts.NotEmpty(req.Header["Authorization"])
|
||||
}
|
||||
@@ -80,6 +80,19 @@ func TestCheckRequest(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
General = HMACAuth{SecretKey: []byte(util.RandStringRunes(256))}
|
||||
|
||||
// 缺少请求头
|
||||
{
|
||||
req, err := http.NewRequest(
|
||||
"POST",
|
||||
"http://127.0.0.1/api/v3/upload",
|
||||
strings.NewReader("I am body."),
|
||||
)
|
||||
asserts.NoError(err)
|
||||
err = CheckRequest(General, req)
|
||||
asserts.Error(err)
|
||||
asserts.Equal(ErrAuthHeaderMissing, err)
|
||||
}
|
||||
|
||||
// 非上传请求 验证成功
|
||||
{
|
||||
req, err := http.NewRequest(
|
||||
@@ -101,7 +114,7 @@ func TestCheckRequest(t *testing.T) {
|
||||
strings.NewReader("I am body."),
|
||||
)
|
||||
asserts.NoError(err)
|
||||
req.Header["X-Policy"] = []string{"I am Policy"}
|
||||
req.Header["X-Cr-Policy"] = []string{"I am Policy"}
|
||||
req = SignRequest(General, req, 0)
|
||||
err = CheckRequest(General, req)
|
||||
asserts.NoError(err)
|
||||
|
||||
+1
-1
@@ -33,7 +33,7 @@ func (auth HMACAuth) Check(body string, sign string) error {
|
||||
signSlice := strings.Split(sign, ":")
|
||||
// 如果未携带expires字段
|
||||
if signSlice[len(signSlice)-1] == "" {
|
||||
return ErrAuthFailed
|
||||
return ErrExpiresMissing
|
||||
}
|
||||
|
||||
// 验证是否过期
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewBalancer(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
a.NotNil(NewBalancer(""))
|
||||
a.IsType(&RoundRobin{}, NewBalancer("RoundRobin"))
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package balancer
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRoundRobin_NextIndex(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
r := &RoundRobin{}
|
||||
total := 5
|
||||
for i := 1; i < total; i++ {
|
||||
a.Equal(i, r.NextIndex(total))
|
||||
}
|
||||
for i := 0; i < total; i++ {
|
||||
a.Equal(i, r.NextIndex(total))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobin_NextPeer(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
r := &RoundRobin{}
|
||||
|
||||
// not slice
|
||||
{
|
||||
err, _ := r.NextPeer("s")
|
||||
a.Equal(ErrInputNotSlice, err)
|
||||
}
|
||||
|
||||
// no nodes
|
||||
{
|
||||
err, _ := r.NextPeer([]string{})
|
||||
a.Equal(ErrNoAvaliableNode, err)
|
||||
}
|
||||
|
||||
// pass
|
||||
{
|
||||
err, res := r.NextPeer([]string{"a"})
|
||||
a.NoError(err)
|
||||
a.Equal("a", res.(string))
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
package slave
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
@@ -51,13 +50,13 @@ type MasterInfo struct {
|
||||
TTL int
|
||||
URL *url.URL
|
||||
// used to invoke aria2 rpc calls
|
||||
Instance cluster.Node
|
||||
Instance Node
|
||||
Client request.Client
|
||||
|
||||
jobTracker map[string]bool
|
||||
}
|
||||
|
||||
func Init() {
|
||||
func InitController() {
|
||||
DefaultController = &slaveController{
|
||||
masters: make(map[string]MasterInfo),
|
||||
}
|
||||
@@ -95,7 +94,7 @@ func (c *slaveController) HandleHeartBeat(req *serializer.NodePingReq) (serializ
|
||||
}, int64(req.CredentialTTL)),
|
||||
),
|
||||
jobTracker: make(map[string]bool),
|
||||
Instance: cluster.NewNodeFromDBModel(&model.Node{
|
||||
Instance: NewNodeFromDBModel(&model.Node{
|
||||
Model: gorm.Model{ID: req.Node.ID},
|
||||
MasterKey: req.Node.MasterKey,
|
||||
Type: model.MasterNodeType,
|
||||
@@ -162,6 +161,7 @@ func (c *slaveController) SubmitTask(id string, job interface{}, hash string, su
|
||||
return nil
|
||||
}
|
||||
|
||||
node.jobTracker[hash] = true
|
||||
submitter(job)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestInitController(t *testing.T) {
|
||||
assert.NotPanics(t, func() {
|
||||
InitController()
|
||||
})
|
||||
}
|
||||
|
||||
func TestSlaveController_HandleHeartBeat(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
c := &slaveController{
|
||||
masters: make(map[string]MasterInfo),
|
||||
}
|
||||
|
||||
// first heart beat
|
||||
{
|
||||
_, err := c.HandleHeartBeat(&serializer.NodePingReq{
|
||||
SiteID: "1",
|
||||
Node: &model.Node{},
|
||||
})
|
||||
a.NoError(err)
|
||||
|
||||
_, err = c.HandleHeartBeat(&serializer.NodePingReq{
|
||||
SiteID: "2",
|
||||
Node: &model.Node{},
|
||||
})
|
||||
a.NoError(err)
|
||||
|
||||
a.Len(c.masters, 2)
|
||||
}
|
||||
|
||||
// second heart beat, no fresh
|
||||
{
|
||||
_, err := c.HandleHeartBeat(&serializer.NodePingReq{
|
||||
SiteID: "1",
|
||||
SiteURL: "http://127.0.0.1",
|
||||
Node: &model.Node{},
|
||||
})
|
||||
a.NoError(err)
|
||||
a.Len(c.masters, 2)
|
||||
a.Empty(c.masters["1"].URL)
|
||||
}
|
||||
|
||||
// second heart beat, fresh
|
||||
{
|
||||
_, err := c.HandleHeartBeat(&serializer.NodePingReq{
|
||||
SiteID: "1",
|
||||
IsUpdate: true,
|
||||
SiteURL: "http://127.0.0.1",
|
||||
Node: &model.Node{},
|
||||
})
|
||||
a.NoError(err)
|
||||
a.Len(c.masters, 2)
|
||||
a.Equal("http://127.0.0.1", c.masters["1"].URL.String())
|
||||
}
|
||||
|
||||
// second heart beat, fresh, url illegal
|
||||
{
|
||||
_, err := c.HandleHeartBeat(&serializer.NodePingReq{
|
||||
SiteID: "1",
|
||||
IsUpdate: true,
|
||||
SiteURL: string([]byte{0x7f}),
|
||||
Node: &model.Node{},
|
||||
})
|
||||
a.Error(err)
|
||||
a.Len(c.masters, 2)
|
||||
a.Equal("http://127.0.0.1", c.masters["1"].URL.String())
|
||||
}
|
||||
}
|
||||
|
||||
type nodeMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (n nodeMock) Init(node *model.Node) {
|
||||
n.Called(node)
|
||||
}
|
||||
|
||||
func (n nodeMock) IsFeatureEnabled(feature string) bool {
|
||||
args := n.Called(feature)
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n nodeMock) SubscribeStatusChange(callback func(isActive bool, id uint)) {
|
||||
n.Called(callback)
|
||||
}
|
||||
|
||||
func (n nodeMock) Ping(req *serializer.NodePingReq) (*serializer.NodePingResp, error) {
|
||||
args := n.Called(req)
|
||||
return args.Get(0).(*serializer.NodePingResp), args.Error(1)
|
||||
}
|
||||
|
||||
func (n nodeMock) IsActive() bool {
|
||||
args := n.Called()
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n nodeMock) GetAria2Instance() common.Aria2 {
|
||||
args := n.Called()
|
||||
return args.Get(0).(common.Aria2)
|
||||
}
|
||||
|
||||
func (n nodeMock) ID() uint {
|
||||
args := n.Called()
|
||||
return args.Get(0).(uint)
|
||||
}
|
||||
|
||||
func (n nodeMock) Kill() {
|
||||
n.Called()
|
||||
}
|
||||
|
||||
func (n nodeMock) IsMater() bool {
|
||||
args := n.Called()
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n nodeMock) MasterAuthInstance() auth.Auth {
|
||||
args := n.Called()
|
||||
return args.Get(0).(auth.Auth)
|
||||
}
|
||||
|
||||
func (n nodeMock) SlaveAuthInstance() auth.Auth {
|
||||
args := n.Called()
|
||||
return args.Get(0).(auth.Auth)
|
||||
}
|
||||
|
||||
func (n nodeMock) DBModel() *model.Node {
|
||||
args := n.Called()
|
||||
return args.Get(0).(*model.Node)
|
||||
}
|
||||
|
||||
func TestSlaveController_GetAria2Instance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mockNode := &nodeMock{}
|
||||
mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Instance: mockNode},
|
||||
},
|
||||
}
|
||||
|
||||
// node node found
|
||||
{
|
||||
res, err := c.GetAria2Instance("2")
|
||||
a.Nil(res)
|
||||
a.Equal(ErrMasterNotFound, err)
|
||||
}
|
||||
|
||||
// node found
|
||||
{
|
||||
res, err := c.GetAria2Instance("1")
|
||||
a.NotNil(res)
|
||||
a.NoError(err)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
type requestMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (r requestMock) Request(method, target string, body io.Reader, opts ...request.Option) *request.Response {
|
||||
return r.Called(method, target, body, opts).Get(0).(*request.Response)
|
||||
}
|
||||
|
||||
func TestSlaveController_SendNotification(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {},
|
||||
},
|
||||
}
|
||||
|
||||
// node not exit
|
||||
{
|
||||
a.Equal(ErrMasterNotFound, c.SendNotification("2", "", mq.Message{}))
|
||||
}
|
||||
|
||||
// gob encode error
|
||||
{
|
||||
type randomType struct{}
|
||||
a.Error(c.SendNotification("1", "", mq.Message{
|
||||
Content: randomType{},
|
||||
}))
|
||||
}
|
||||
|
||||
// return none 200
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s1", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{StatusCode: http.StatusConflict},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
a.Error(c.SendNotification("1", "s1", mq.Message{}))
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s2", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
a.Equal(1, c.SendNotification("1", "s2", mq.Message{}).(serializer.AppError).Code)
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s3", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":0}")),
|
||||
},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
a.NoError(c.SendNotification("1", "s3", mq.Message{}))
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveController_SubmitTask(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {
|
||||
jobTracker: map[string]bool{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// node not exit
|
||||
{
|
||||
a.Equal(ErrMasterNotFound, c.SubmitTask("2", "", "", nil))
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
submitted := false
|
||||
a.NoError(c.SubmitTask("1", "", "hash", func(i interface{}) {
|
||||
submitted = true
|
||||
}))
|
||||
a.True(submitted)
|
||||
}
|
||||
|
||||
// job already submitted
|
||||
{
|
||||
submitted := false
|
||||
a.NoError(c.SubmitTask("1", "", "hash", func(i interface{}) {
|
||||
submitted = true
|
||||
}))
|
||||
a.False(submitted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveController_GetMasterInfo(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {},
|
||||
},
|
||||
}
|
||||
|
||||
// node not exit
|
||||
{
|
||||
res, err := c.GetMasterInfo("2")
|
||||
a.Equal(ErrMasterNotFound, err)
|
||||
a.Nil(res)
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
res, err := c.GetMasterInfo("1")
|
||||
a.NoError(err)
|
||||
a.NotNil(res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveController_GetOneDriveToken(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {},
|
||||
},
|
||||
}
|
||||
|
||||
// node not exit
|
||||
{
|
||||
res, err := c.GetOneDriveToken("2", 1)
|
||||
a.Equal(ErrMasterNotFound, err)
|
||||
a.Empty(res)
|
||||
}
|
||||
|
||||
// return none 200
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{StatusCode: http.StatusConflict},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
res, err := c.GetOneDriveToken("1", 1)
|
||||
a.Error(err)
|
||||
a.Empty(res)
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
res, err := c.GetOneDriveToken("1", 1)
|
||||
a.Equal(1, err.(serializer.AppError).Code)
|
||||
a.Empty(res)
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// success
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"expected\"}")),
|
||||
},
|
||||
})
|
||||
c := &slaveController{
|
||||
masters: map[string]MasterInfo{
|
||||
"1": {Client: mockRequest},
|
||||
},
|
||||
}
|
||||
res, err := c.GetOneDriveToken("1", 1)
|
||||
a.NoError(err)
|
||||
a.Equal("expected", res)
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
}
|
||||
@@ -1,8 +1,12 @@
|
||||
package cluster
|
||||
|
||||
import "errors"
|
||||
import (
|
||||
"errors"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrFeatureNotExist = errors.New("No nodes in nodepool match the feature specificed")
|
||||
ErrIlegalPath = errors.New("path out of boundary of setting temp folder")
|
||||
ErrMasterNotFound = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
|
||||
)
|
||||
|
||||
+15
-8
@@ -20,7 +20,10 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
const deleteTempFileDuration = 60 * time.Second
|
||||
const (
|
||||
deleteTempFileDuration = 60 * time.Second
|
||||
statusRetryDuration = 10 * time.Second
|
||||
)
|
||||
|
||||
type MasterNode struct {
|
||||
Model *model.Node
|
||||
@@ -33,8 +36,10 @@ type rpcService struct {
|
||||
Caller rpc.Client
|
||||
Initialized bool
|
||||
|
||||
parent *MasterNode
|
||||
options *clientOptions
|
||||
retryDuration time.Duration
|
||||
deletePaddingDuration time.Duration
|
||||
parent *MasterNode
|
||||
options *clientOptions
|
||||
}
|
||||
|
||||
type clientOptions struct {
|
||||
@@ -46,6 +51,8 @@ func (node *MasterNode) Init(nodeModel *model.Node) {
|
||||
node.lock.Lock()
|
||||
node.Model = nodeModel
|
||||
node.aria2RPC.parent = node
|
||||
node.aria2RPC.retryDuration = statusRetryDuration
|
||||
node.aria2RPC.deletePaddingDuration = deleteTempFileDuration
|
||||
node.lock.Unlock()
|
||||
|
||||
node.lock.RLock()
|
||||
@@ -214,8 +221,8 @@ func (r *rpcService) Status(task *model.Download) (rpc.StatusInfo, error) {
|
||||
res, err := r.Caller.TellStatus(task.GID)
|
||||
if err != nil {
|
||||
// 失败后重试
|
||||
util.Log().Debug("无法获取离线下载状态,%s,10秒钟后重试", err)
|
||||
time.Sleep(time.Duration(10) * time.Second)
|
||||
util.Log().Debug("无法获取离线下载状态,%s,稍后重试", err)
|
||||
time.Sleep(r.retryDuration)
|
||||
res, err = r.Caller.TellStatus(task.GID)
|
||||
}
|
||||
|
||||
@@ -253,13 +260,13 @@ func (s *rpcService) DeleteTempFile(task *model.Download) error {
|
||||
defer s.parent.lock.RUnlock()
|
||||
|
||||
// 避免被aria2占用,异步执行删除
|
||||
go func(src string) {
|
||||
time.Sleep(deleteTempFileDuration)
|
||||
go func(d time.Duration, src string) {
|
||||
time.Sleep(d)
|
||||
err := os.RemoveAll(src)
|
||||
if err != nil {
|
||||
util.Log().Warning("无法删除离线下载临时目录[%s], %s", src, err)
|
||||
}
|
||||
}(task.Parent)
|
||||
}(s.deletePaddingDuration, task.Parent)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"context"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestMasterNode_Init(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{}
|
||||
m.Init(&model.Node{Status: model.NodeSuspend})
|
||||
a.Equal(model.NodeSuspend, m.DBModel().Status)
|
||||
m.Init(&model.Node{Aria2Enabled: true})
|
||||
}
|
||||
|
||||
func TestMasterNode_DummyMethods(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
m.Model.ID = 5
|
||||
a.Equal(m.Model.ID, m.ID())
|
||||
|
||||
res, err := m.Ping(&serializer.NodePingReq{})
|
||||
a.NoError(err)
|
||||
a.NotNil(res)
|
||||
|
||||
a.True(m.IsActive())
|
||||
a.True(m.IsMater())
|
||||
|
||||
m.SubscribeStatusChange(func(isActive bool, id uint) {})
|
||||
}
|
||||
|
||||
func TestMasterNode_IsFeatureEnabled(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
a.False(m.IsFeatureEnabled("aria2"))
|
||||
a.False(m.IsFeatureEnabled("random"))
|
||||
m.Model.Aria2Enabled = true
|
||||
a.True(m.IsFeatureEnabled("aria2"))
|
||||
}
|
||||
|
||||
func TestMasterNode_AuthInstance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
a.NotNil(m.MasterAuthInstance())
|
||||
a.NotNil(m.SlaveAuthInstance())
|
||||
}
|
||||
|
||||
func TestMasterNode_Kill(t *testing.T) {
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
m.Kill()
|
||||
|
||||
caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
|
||||
m.aria2RPC.Caller = caller
|
||||
m.Kill()
|
||||
}
|
||||
|
||||
func TestMasterNode_GetAria2Instance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{},
|
||||
aria2RPC: rpcService{},
|
||||
}
|
||||
|
||||
m.aria2RPC.parent = m
|
||||
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
m.Model.Aria2Enabled = true
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
m.aria2RPC.Initialized = true
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
}
|
||||
|
||||
func TestRpcService_Init(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{
|
||||
Aria2OptionsSerialized: model.Aria2Option{
|
||||
Options: "{",
|
||||
},
|
||||
},
|
||||
aria2RPC: rpcService{},
|
||||
}
|
||||
m.aria2RPC.parent = m
|
||||
|
||||
// failed to decode address
|
||||
{
|
||||
m.Model.Aria2OptionsSerialized.Server = string([]byte{0x7f})
|
||||
a.Error(m.aria2RPC.Init())
|
||||
}
|
||||
|
||||
// failed to decode options
|
||||
{
|
||||
m.Model.Aria2OptionsSerialized.Server = ""
|
||||
a.Error(m.aria2RPC.Init())
|
||||
}
|
||||
|
||||
// failed to initialized
|
||||
{
|
||||
m.Model.Aria2OptionsSerialized.Server = ""
|
||||
m.Model.Aria2OptionsSerialized.Options = "{}"
|
||||
caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
|
||||
m.aria2RPC.Caller = caller
|
||||
a.Error(m.aria2RPC.Init())
|
||||
a.False(m.aria2RPC.Initialized)
|
||||
}
|
||||
}
|
||||
|
||||
func getTestRPCNode() *MasterNode {
|
||||
m := &MasterNode{
|
||||
Model: &model.Node{
|
||||
Aria2OptionsSerialized: model.Aria2Option{},
|
||||
},
|
||||
aria2RPC: rpcService{
|
||||
options: &clientOptions{
|
||||
Options: map[string]interface{}{"1": "1"},
|
||||
},
|
||||
},
|
||||
}
|
||||
m.aria2RPC.parent = m
|
||||
caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
|
||||
m.aria2RPC.Caller = caller
|
||||
return m
|
||||
}
|
||||
|
||||
func TestRpcService_CreateTask(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNode()
|
||||
|
||||
res, err := m.aria2RPC.CreateTask(&model.Download{}, map[string]interface{}{"1": "1"})
|
||||
a.Error(err)
|
||||
a.Empty(res)
|
||||
}
|
||||
|
||||
func TestRpcService_Status(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNode()
|
||||
|
||||
res, err := m.aria2RPC.Status(&model.Download{})
|
||||
a.Error(err)
|
||||
a.Empty(res)
|
||||
}
|
||||
|
||||
func TestRpcService_Cancel(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNode()
|
||||
|
||||
a.Error(m.aria2RPC.Cancel(&model.Download{}))
|
||||
}
|
||||
|
||||
func TestRpcService_Select(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNode()
|
||||
|
||||
a.NotNil(m.aria2RPC.GetConfig())
|
||||
a.Error(m.aria2RPC.Select(&model.Download{}, []int{1, 2, 3}))
|
||||
}
|
||||
|
||||
func TestRpcService_DeleteTempFile(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNode()
|
||||
fdName := "TestRpcService_DeleteTempFile"
|
||||
a.NoError(os.Mkdir(fdName, 0644))
|
||||
|
||||
a.NoError(m.aria2RPC.DeleteTempFile(&model.Download{Parent: fdName}))
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
a.False(util.Exists(fdName))
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewNodeFromDBModel(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
a.IsType(&SlaveNode{}, NewNodeFromDBModel(&model.Node{
|
||||
Type: model.SlaveNodeType,
|
||||
}))
|
||||
a.IsType(&MasterNode{}, NewNodeFromDBModel(&model.Node{
|
||||
Type: model.MasterNodeType,
|
||||
}))
|
||||
}
|
||||
+20
-9
@@ -39,14 +39,22 @@ type NodePool struct {
|
||||
|
||||
// Init 初始化从机节点池
|
||||
func Init() {
|
||||
Default = &NodePool{
|
||||
featureMap: make(map[string][]Node),
|
||||
}
|
||||
Default = &NodePool{}
|
||||
Default.Init()
|
||||
if err := Default.initFromDB(); err != nil {
|
||||
util.Log().Warning("节点池初始化失败, %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (pool *NodePool) Init() {
|
||||
pool.lock.Lock()
|
||||
defer pool.lock.Unlock()
|
||||
|
||||
pool.featureMap = make(map[string][]Node)
|
||||
pool.active = make(map[uint]Node)
|
||||
pool.inactive = make(map[uint]Node)
|
||||
}
|
||||
|
||||
func (pool *NodePool) buildIndexMap() {
|
||||
pool.lock.Lock()
|
||||
for _, feature := range featureGroup {
|
||||
@@ -76,14 +84,19 @@ func (pool *NodePool) GetNodeByID(id uint) Node {
|
||||
|
||||
func (pool *NodePool) nodeStatusChange(isActive bool, id uint) {
|
||||
util.Log().Debug("从机节点 [ID=%d] 状态变更 [Active=%t]", id, isActive)
|
||||
var node Node
|
||||
pool.lock.Lock()
|
||||
if isActive {
|
||||
node := pool.inactive[id]
|
||||
if n, ok := pool.inactive[id]; ok {
|
||||
node = n
|
||||
delete(pool.inactive, id)
|
||||
} else {
|
||||
node = pool.active[id]
|
||||
delete(pool.active, id)
|
||||
}
|
||||
|
||||
if isActive {
|
||||
pool.active[id] = node
|
||||
} else {
|
||||
node := pool.active[id]
|
||||
delete(pool.active, id)
|
||||
pool.inactive[id] = node
|
||||
}
|
||||
pool.lock.Unlock()
|
||||
@@ -98,8 +111,6 @@ func (pool *NodePool) initFromDB() error {
|
||||
}
|
||||
|
||||
pool.lock.Lock()
|
||||
pool.active = make(map[uint]Node)
|
||||
pool.inactive = make(map[uint]Node)
|
||||
for i := 0; i < len(nodes); i++ {
|
||||
pool.add(&nodes[i])
|
||||
}
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
|
||||
"github.com/jinzhu/gorm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
var mock sqlmock.Sqlmock
|
||||
|
||||
// TestMain 初始化数据库Mock
|
||||
func TestMain(m *testing.M) {
|
||||
var db *sql.DB
|
||||
var err error
|
||||
db, mock, err = sqlmock.New()
|
||||
if err != nil {
|
||||
panic("An error was not expected when opening a stub database connection")
|
||||
}
|
||||
model.DB, _ = gorm.Open("mysql", db)
|
||||
defer db.Close()
|
||||
m.Run()
|
||||
}
|
||||
|
||||
func TestInitFailed(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnError(errors.New("error"))
|
||||
Init()
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestInitSuccess(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "aria2_enabled", "type"}).AddRow(1, true, model.MasterNodeType))
|
||||
Init()
|
||||
a.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestNodePool_GetNodeByID(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
p := &NodePool{}
|
||||
p.Init()
|
||||
mockNode := &nodeMock{}
|
||||
|
||||
// inactive
|
||||
{
|
||||
p.inactive[1] = mockNode
|
||||
a.Equal(mockNode, p.GetNodeByID(1))
|
||||
}
|
||||
|
||||
// active
|
||||
{
|
||||
delete(p.inactive, 1)
|
||||
p.active[1] = mockNode
|
||||
a.Equal(mockNode, p.GetNodeByID(1))
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodePool_NodeStatusChange(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
p := &NodePool{}
|
||||
n := &MasterNode{Model: &model.Node{}}
|
||||
p.Init()
|
||||
p.inactive[1] = n
|
||||
|
||||
p.nodeStatusChange(true, 1)
|
||||
a.Len(p.inactive, 0)
|
||||
a.Equal(n, p.active[1])
|
||||
|
||||
p.nodeStatusChange(false, 1)
|
||||
a.Len(p.active, 0)
|
||||
a.Equal(n, p.inactive[1])
|
||||
|
||||
p.nodeStatusChange(false, 1)
|
||||
a.Len(p.active, 0)
|
||||
a.Equal(n, p.inactive[1])
|
||||
}
|
||||
|
||||
func TestNodePool_Add(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
p := &NodePool{}
|
||||
p.Init()
|
||||
|
||||
// new node
|
||||
{
|
||||
p.Add(&model.Node{})
|
||||
a.Len(p.active, 1)
|
||||
}
|
||||
|
||||
// old node
|
||||
{
|
||||
p.inactive[0] = p.active[0]
|
||||
delete(p.active, 0)
|
||||
p.Add(&model.Node{})
|
||||
a.Len(p.active, 0)
|
||||
a.Len(p.inactive, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodePool_Delete(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
p := &NodePool{}
|
||||
p.Init()
|
||||
|
||||
// active
|
||||
{
|
||||
mockNode := &nodeMock{}
|
||||
mockNode.On("Kill")
|
||||
p.active[0] = mockNode
|
||||
p.Delete(0)
|
||||
a.Len(p.active, 0)
|
||||
a.Len(p.inactive, 0)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
|
||||
p.Init()
|
||||
|
||||
// inactive
|
||||
{
|
||||
mockNode := &nodeMock{}
|
||||
mockNode.On("Kill")
|
||||
p.inactive[0] = mockNode
|
||||
p.Delete(0)
|
||||
a.Len(p.active, 0)
|
||||
a.Len(p.inactive, 0)
|
||||
mockNode.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNodePool_BalanceNodeByFeature(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
p := &NodePool{}
|
||||
p.Init()
|
||||
|
||||
// success
|
||||
{
|
||||
p.featureMap["test"] = []Node{&MasterNode{}}
|
||||
err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
|
||||
a.NoError(err)
|
||||
a.Equal(p.featureMap["test"][0], res)
|
||||
}
|
||||
|
||||
// NoNodes
|
||||
{
|
||||
p.featureMap["test"] = []Node{}
|
||||
err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
|
||||
a.Error(err)
|
||||
a.Nil(res)
|
||||
}
|
||||
|
||||
// No match feature
|
||||
{
|
||||
err, res := p.BalanceNodeByFeature("test2", balancer.NewBalancer("round-robin"))
|
||||
a.Error(err)
|
||||
a.Nil(res)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/stretchr/testify/assert"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSlaveNode_InitAndKill(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
n := &SlaveNode{
|
||||
callback: func(b bool, u uint) {
|
||||
|
||||
},
|
||||
}
|
||||
|
||||
a.NotPanics(func() {
|
||||
n.Init(&model.Node{})
|
||||
time.Sleep(time.Millisecond * 500)
|
||||
n.Init(&model.Node{})
|
||||
n.Kill()
|
||||
})
|
||||
}
|
||||
|
||||
func TestSlaveNode_DummyMethods(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
m.Model.ID = 5
|
||||
a.Equal(m.Model.ID, m.ID())
|
||||
a.Equal(m.Model.ID, m.DBModel().ID)
|
||||
|
||||
a.False(m.IsActive())
|
||||
a.False(m.IsMater())
|
||||
|
||||
m.SubscribeStatusChange(func(isActive bool, id uint) {})
|
||||
}
|
||||
|
||||
func TestSlaveNode_IsFeatureEnabled(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
a.False(m.IsFeatureEnabled("aria2"))
|
||||
a.False(m.IsFeatureEnabled("random"))
|
||||
m.Model.Aria2Enabled = true
|
||||
a.True(m.IsFeatureEnabled("aria2"))
|
||||
}
|
||||
|
||||
func TestSlaveNode_Ping(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
// master return error code
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.Ping(&serializer.NodePingReq{})
|
||||
a.Error(err)
|
||||
a.Nil(res)
|
||||
a.Equal(1, err.(serializer.AppError).Code)
|
||||
}
|
||||
|
||||
// return unexpected json
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"233\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.Ping(&serializer.NodePingReq{})
|
||||
a.Error(err)
|
||||
a.Nil(res)
|
||||
}
|
||||
|
||||
// return success
|
||||
{
|
||||
mockRequest := &requestMock{}
|
||||
mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"{}\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.Ping(&serializer.NodePingReq{})
|
||||
a.NoError(err)
|
||||
a.NotNil(res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveNode_GetAria2Instance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
m.Model.Aria2Enabled = true
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
a.NotNil(m.GetAria2Instance())
|
||||
}
|
||||
|
||||
func TestSlaveNode_StartPingLoop(t *testing.T) {
|
||||
callbackCount := 0
|
||||
finishedChan := make(chan struct{})
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m := &SlaveNode{
|
||||
Active: true,
|
||||
Model: &model.Node{},
|
||||
callback: func(b bool, u uint) {
|
||||
callbackCount++
|
||||
if callbackCount == 2 {
|
||||
close(finishedChan)
|
||||
}
|
||||
if callbackCount == 1 {
|
||||
mockRequest.AssertExpectations(t)
|
||||
mockRequest = requestMock{}
|
||||
mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"{}\"}")),
|
||||
},
|
||||
})
|
||||
}
|
||||
},
|
||||
}
|
||||
cache.Set("setting_slave_ping_interval", "0", 0)
|
||||
cache.Set("setting_slave_recover_interval", "0", 0)
|
||||
cache.Set("setting_slave_node_retry", "1", 0)
|
||||
|
||||
m.caller.Client = &mockRequest
|
||||
go func() {
|
||||
select {
|
||||
case <-finishedChan:
|
||||
m.Kill()
|
||||
}
|
||||
}()
|
||||
m.StartPingLoop()
|
||||
mockRequest.AssertExpectations(t)
|
||||
}
|
||||
|
||||
func TestSlaveNode_AuthInstance(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
a.NotNil(m.MasterAuthInstance())
|
||||
a.NotNil(m.SlaveAuthInstance())
|
||||
}
|
||||
|
||||
func TestSlaveNode_ChangeStatus(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
isActive := false
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
callback: func(b bool, u uint) {
|
||||
isActive = b
|
||||
},
|
||||
}
|
||||
|
||||
a.NotPanics(func() {
|
||||
m.changeStatus(false)
|
||||
})
|
||||
m.changeStatus(true)
|
||||
a.True(isActive)
|
||||
}
|
||||
|
||||
func getTestRPCNodeSlave() *SlaveNode {
|
||||
m := &SlaveNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
m.caller.parent = m
|
||||
return m
|
||||
}
|
||||
|
||||
func TestSlaveCaller_CreateTask(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNodeSlave()
|
||||
|
||||
// master return 404
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.CreateTask(&model.Download{}, nil)
|
||||
a.Empty(res)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.CreateTask(&model.Download{}, nil)
|
||||
a.Empty(res)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return success
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.CreateTask(&model.Download{}, nil)
|
||||
a.Equal("res", res)
|
||||
a.NoError(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveCaller_Status(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNodeSlave()
|
||||
|
||||
// master return 404
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.Status(&model.Download{})
|
||||
a.Empty(res.Status)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.Status(&model.Download{})
|
||||
a.Empty(res.Status)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return success
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"re456456s\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
res, err := m.caller.Status(&model.Download{})
|
||||
a.Empty(res.Status)
|
||||
a.NoError(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveCaller_Cancel(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNodeSlave()
|
||||
|
||||
// master return 404
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Cancel(&model.Download{})
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Cancel(&model.Download{})
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return success
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Cancel(&model.Download{})
|
||||
a.NoError(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveCaller_Select(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNodeSlave()
|
||||
m.caller.Init()
|
||||
m.caller.GetConfig()
|
||||
|
||||
// master return 404
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Select(&model.Download{}, nil)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Select(&model.Download{}, nil)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return success
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.Select(&model.Download{}, nil)
|
||||
a.NoError(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSlaveCaller_DeleteTempFile(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
m := getTestRPCNodeSlave()
|
||||
m.caller.Init()
|
||||
m.caller.GetConfig()
|
||||
|
||||
// master return 404
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 404,
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.DeleteTempFile(&model.Download{})
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return error
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.DeleteTempFile(&model.Download{})
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
// master return success
|
||||
{
|
||||
mockRequest := requestMock{}
|
||||
mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
|
||||
Response: &http.Response{
|
||||
StatusCode: 200,
|
||||
Body: ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
|
||||
},
|
||||
})
|
||||
m.caller.Client = mockRequest
|
||||
err := m.caller.DeleteTempFile(&model.Download{})
|
||||
a.NoError(err)
|
||||
}
|
||||
}
|
||||
+7
-3
@@ -70,9 +70,13 @@ type redis struct {
|
||||
|
||||
// 缩略图 配置
|
||||
type thumb struct {
|
||||
MaxWidth uint
|
||||
MaxHeight uint
|
||||
FileSuffix string `validate:"min=1"`
|
||||
MaxWidth uint
|
||||
MaxHeight uint
|
||||
FileSuffix string `validate:"min=1"`
|
||||
MaxTaskCount int
|
||||
EncodeMethod string `validate:"eq=jpg|eq=png"`
|
||||
EncodeQuality int `validate:"gte=1,lte=100"`
|
||||
GCAfterGen bool
|
||||
}
|
||||
|
||||
// 跨域配置
|
||||
|
||||
@@ -44,16 +44,20 @@ var CaptchaConfig = &captcha{
|
||||
var CORSConfig = &cors{
|
||||
AllowOrigins: []string{"UNSET"},
|
||||
AllowMethods: []string{"PUT", "POST", "GET", "OPTIONS"},
|
||||
AllowHeaders: []string{"Cookie", "X-Policy", "Authorization", "Content-Length", "Content-Type", "X-Path", "X-FileName"},
|
||||
AllowHeaders: []string{"Cookie", "X-Cr-Policy", "Authorization", "Content-Length", "Content-Type", "X-Cr-Path", "X-Cr-FileName"},
|
||||
AllowCredentials: false,
|
||||
ExposeHeaders: nil,
|
||||
}
|
||||
|
||||
// ThumbConfig 缩略图配置
|
||||
var ThumbConfig = &thumb{
|
||||
MaxWidth: 400,
|
||||
MaxHeight: 300,
|
||||
FileSuffix: "._thumb",
|
||||
MaxWidth: 400,
|
||||
MaxHeight: 300,
|
||||
FileSuffix: "._thumb",
|
||||
MaxTaskCount: -1,
|
||||
EncodeMethod: "jpg",
|
||||
GCAfterGen: false,
|
||||
EncodeQuality: 85,
|
||||
}
|
||||
|
||||
// SlaveConfig 从机配置
|
||||
|
||||
+3
-3
@@ -1,13 +1,13 @@
|
||||
package conf
|
||||
|
||||
// BackendVersion 当前后端版本号
|
||||
var BackendVersion = "3.4.0-beta1"
|
||||
var BackendVersion = "3.4.3"
|
||||
|
||||
// RequiredDBVersion 与当前版本匹配的数据库版本
|
||||
var RequiredDBVersion = "3.4.0-beta1"
|
||||
var RequiredDBVersion = "3.4.0"
|
||||
|
||||
// RequiredStaticVersion 与当前版本匹配的静态资源版本
|
||||
var RequiredStaticVersion = "3.4.0-beta1"
|
||||
var RequiredStaticVersion = "3.4.3"
|
||||
|
||||
// IsPro 是否为Pro版本
|
||||
var IsPro = "false"
|
||||
|
||||
@@ -190,7 +190,7 @@ func TestHandler_Source(t *testing.T) {
|
||||
|
||||
// 设定了CDN,解析失败
|
||||
{
|
||||
handler.Policy.BaseURL = string(0x7f)
|
||||
handler.Policy.BaseURL = string([]byte{0x7f})
|
||||
file := model.File{
|
||||
Model: gorm.Model{
|
||||
ID: 1,
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
@@ -573,7 +574,7 @@ func sysError(err error) *RespError {
|
||||
|
||||
func (client *Client) request(ctx context.Context, method string, url string, body io.Reader, option ...request.Option) (string, *RespError) {
|
||||
// 获取凭证
|
||||
err := client.UpdateCredential(ctx)
|
||||
err := client.UpdateCredential(ctx, conf.SystemConfig.Mode == "slave")
|
||||
if err != nil {
|
||||
return "", sysError(err)
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package onedrive
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
@@ -28,7 +29,8 @@ type Client struct {
|
||||
ClientSecret string
|
||||
Redirect string
|
||||
|
||||
Request request.Client
|
||||
Request request.Client
|
||||
ClusterController cluster.Controller
|
||||
}
|
||||
|
||||
// Endpoints OneDrive客户端相关设置
|
||||
@@ -51,11 +53,12 @@ func NewClient(policy *model.Policy) (*Client, error) {
|
||||
Credential: &Credential{
|
||||
RefreshToken: policy.AccessKey,
|
||||
},
|
||||
Policy: policy,
|
||||
ClientID: policy.BucketName,
|
||||
ClientSecret: policy.SecretKey,
|
||||
Redirect: policy.OptionsSerialized.OdRedirect,
|
||||
Request: request.NewClient(),
|
||||
Policy: policy,
|
||||
ClientID: policy.BucketName,
|
||||
ClientSecret: policy.SecretKey,
|
||||
Redirect: policy.OptionsSerialized.OdRedirect,
|
||||
Request: request.NewClient(),
|
||||
ClusterController: cluster.DefaultController,
|
||||
}
|
||||
|
||||
if client.Endpoints.DriverResource == "" {
|
||||
|
||||
@@ -10,9 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/slave"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
)
|
||||
|
||||
@@ -125,8 +123,8 @@ func (client *Client) ObtainToken(ctx context.Context, opts ...Option) (*Credent
|
||||
}
|
||||
|
||||
// UpdateCredential 更新凭证,并检查有效期
|
||||
func (client *Client) UpdateCredential(ctx context.Context) error {
|
||||
if conf.SystemConfig.Mode == "slave" {
|
||||
func (client *Client) UpdateCredential(ctx context.Context, isSlave bool) error {
|
||||
if isSlave {
|
||||
return client.fetchCredentialFromMaster(ctx)
|
||||
}
|
||||
|
||||
@@ -179,7 +177,7 @@ func (client *Client) UpdateCredential(ctx context.Context) error {
|
||||
|
||||
// UpdateCredential 更新凭证,并检查有效期
|
||||
func (client *Client) fetchCredentialFromMaster(ctx context.Context) error {
|
||||
res, err := slave.DefaultController.GetOneDriveToken(client.Policy.MasterID, client.Policy.ID)
|
||||
res, err := client.ClusterController.GetOneDriveToken(client.Policy.MasterID, client.Policy.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mocks/controllermock"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
@@ -269,10 +270,10 @@ func TestClient_UpdateCredential(t *testing.T) {
|
||||
|
||||
// 无有效的RefreshToken
|
||||
{
|
||||
err := client.UpdateCredential(context.Background())
|
||||
err := client.UpdateCredential(context.Background(), false)
|
||||
asserts.Equal(ErrInvalidRefreshToken, err)
|
||||
client.Credential = nil
|
||||
err = client.UpdateCredential(context.Background())
|
||||
err = client.UpdateCredential(context.Background(), false)
|
||||
asserts.Equal(ErrInvalidRefreshToken, err)
|
||||
}
|
||||
|
||||
@@ -299,7 +300,7 @@ func TestClient_UpdateCredential(t *testing.T) {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
err := client.UpdateCredential(context.Background())
|
||||
err := client.UpdateCredential(context.Background(), false)
|
||||
clientMock.AssertExpectations(t)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.NoError(err)
|
||||
@@ -331,7 +332,7 @@ func TestClient_UpdateCredential(t *testing.T) {
|
||||
client.Credential = &Credential{
|
||||
RefreshToken: "old_refresh_token",
|
||||
}
|
||||
err := client.UpdateCredential(context.Background())
|
||||
err := client.UpdateCredential(context.Background(), false)
|
||||
clientMock.AssertExpectations(t)
|
||||
asserts.Error(err)
|
||||
}
|
||||
@@ -346,7 +347,7 @@ func TestClient_UpdateCredential(t *testing.T) {
|
||||
client.Credential = &Credential{
|
||||
RefreshToken: "old_refresh_token",
|
||||
}
|
||||
err := client.UpdateCredential(context.Background())
|
||||
err := client.UpdateCredential(context.Background(), false)
|
||||
asserts.NoError(err)
|
||||
asserts.Equal("AccessToken", client.Credential.AccessToken)
|
||||
asserts.Equal("RefreshToken", client.Credential.RefreshToken)
|
||||
@@ -359,8 +360,27 @@ func TestClient_UpdateCredential(t *testing.T) {
|
||||
AccessToken: "AccessToken2",
|
||||
ExpiresIn: time.Now().Add(time.Duration(10) * time.Second).Unix(),
|
||||
}
|
||||
err := client.UpdateCredential(context.Background())
|
||||
err := client.UpdateCredential(context.Background(), false)
|
||||
asserts.NoError(err)
|
||||
asserts.Equal("AccessToken2", client.Credential.AccessToken)
|
||||
}
|
||||
|
||||
// slave failed
|
||||
{
|
||||
mockController := &controllermock.SlaveControllerMock{}
|
||||
mockController.On("GetOneDriveToken", testMock.Anything, testMock.Anything).Return("", errors.New("error"))
|
||||
client.ClusterController = mockController
|
||||
err := client.UpdateCredential(context.Background(), true)
|
||||
asserts.Error(err)
|
||||
}
|
||||
|
||||
// slave success
|
||||
{
|
||||
mockController := &controllermock.SlaveControllerMock{}
|
||||
mockController.On("GetOneDriveToken", testMock.Anything, testMock.Anything).Return("AccessToken3", nil)
|
||||
client.ClusterController = mockController
|
||||
err := client.UpdateCredential(context.Background(), true)
|
||||
asserts.NoError(err)
|
||||
asserts.Equal("AccessToken3", client.Credential.AccessToken)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,10 +153,7 @@ func (handler Driver) Put(ctx context.Context, file io.ReadCloser, dst string, s
|
||||
}
|
||||
|
||||
// 对文件名进行URLEncode
|
||||
fileName, err := url.QueryUnescape(path.Base(dst))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fileName := url.QueryEscape(path.Base(dst))
|
||||
|
||||
// 决定是否要禁用文件覆盖
|
||||
overwrite := "true"
|
||||
@@ -170,9 +167,9 @@ func (handler Driver) Put(ctx context.Context, file io.ReadCloser, dst string, s
|
||||
handler.Policy.GetUploadURL(),
|
||||
file,
|
||||
request.WithHeader(map[string][]string{
|
||||
"X-Policy": {credential.Policy},
|
||||
"X-FileName": {fileName},
|
||||
"X-Overwrite": {overwrite},
|
||||
"X-Cr-Policy": {credential.Policy},
|
||||
"X-Cr-FileName": {fileName},
|
||||
"X-Cr-Overwrite": {overwrite},
|
||||
}),
|
||||
request.WithContentLength(int64(size)),
|
||||
request.WithTimeout(time.Duration(0)),
|
||||
@@ -294,7 +291,7 @@ func (handler Driver) Source(
|
||||
sourcePath := base64.RawURLEncoding.EncodeToString([]byte(path))
|
||||
signedURI, err = auth.SignURI(
|
||||
handler.AuthInstance,
|
||||
fmt.Sprintf("%s/%d/%s/%s", controller, speed, sourcePath, fileName),
|
||||
fmt.Sprintf("%s/%d/%s/%s", controller, speed, sourcePath, url.PathEscape(fileName)),
|
||||
ttl,
|
||||
)
|
||||
|
||||
@@ -335,8 +332,8 @@ func (handler Driver) getUploadCredential(ctx context.Context, policy serializer
|
||||
// 签名上传策略
|
||||
uploadRequest, _ := http.NewRequest("POST", "/api/v3/slave/upload", nil)
|
||||
uploadRequest.Header = map[string][]string{
|
||||
"X-Policy": {policyEncoded},
|
||||
"X-Overwrite": {"false"},
|
||||
"X-Cr-Policy": {policyEncoded},
|
||||
"X-Cr-Overwrite": {"false"},
|
||||
}
|
||||
auth.SignRequest(handler.AuthInstance, uploadRequest, TTL)
|
||||
|
||||
|
||||
@@ -105,7 +105,7 @@ func TestHandler_Source(t *testing.T) {
|
||||
// 解析失败 自定义CDN
|
||||
{
|
||||
handler := Driver{
|
||||
Policy: &model.Policy{Server: "/", BaseURL: string(0x7f)},
|
||||
Policy: &model.Policy{Server: "/", BaseURL: string([]byte{0x7f})},
|
||||
AuthInstance: auth.HMACAuth{},
|
||||
}
|
||||
file := model.File{
|
||||
|
||||
+10
-10
@@ -228,25 +228,25 @@ func TestFileSystem_deleteGroupedFile(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
// 全部失败
|
||||
// 全部不存在
|
||||
{
|
||||
failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
|
||||
asserts.Equal(map[uint][]string{
|
||||
1: {"1_1.txt", "1_2.txt"},
|
||||
2: {"2_1.txt", "2_2.txt"},
|
||||
3: {"3_1.txt"},
|
||||
1: {},
|
||||
2: {},
|
||||
3: {},
|
||||
}, failed)
|
||||
}
|
||||
// 部分失败
|
||||
// 部分不存在
|
||||
{
|
||||
file, err := os.Create(util.RelativePath("1_1.txt"))
|
||||
asserts.NoError(err)
|
||||
_ = file.Close()
|
||||
failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
|
||||
asserts.Equal(map[uint][]string{
|
||||
1: {"1_2.txt"},
|
||||
2: {"2_1.txt", "2_2.txt"},
|
||||
3: {"3_1.txt"},
|
||||
1: {},
|
||||
2: {},
|
||||
3: {},
|
||||
}, failed)
|
||||
}
|
||||
// 部分失败,包含整组未知存储策略导致的失败
|
||||
@@ -259,9 +259,9 @@ func TestFileSystem_deleteGroupedFile(t *testing.T) {
|
||||
files[3].Policy.Type = "unknown"
|
||||
failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
|
||||
asserts.Equal(map[uint][]string{
|
||||
1: {"1_2.txt"},
|
||||
1: {},
|
||||
2: {"2_1.txt", "2_2.txt"},
|
||||
3: {"3_1.txt"},
|
||||
3: {},
|
||||
}, failed)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -247,9 +247,6 @@ func (fs *FileSystem) SwitchToSlaveHandler(node cluster.Node) {
|
||||
// SwitchToShadowHandler 将负责上传的 Handler 切换为从机节点转存使用的影子处理器
|
||||
func (fs *FileSystem) SwitchToShadowHandler(master cluster.Node, masterURL, masterID string) {
|
||||
switch fs.Policy.Type {
|
||||
case "remote":
|
||||
fs.Policy.Type = "local"
|
||||
fs.DispatchHandler()
|
||||
case "local":
|
||||
fs.Policy.Type = "remote"
|
||||
fs.Policy.Server = masterURL
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
package filesystem
|
||||
|
||||
import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/shadow/masterinslave"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/shadow/slaveinmaster"
|
||||
"net/http/httptest"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
@@ -104,6 +107,10 @@ func TestDispatchHandler(t *testing.T) {
|
||||
err = fs.DispatchHandler()
|
||||
asserts.NoError(err)
|
||||
|
||||
fs.Policy = &model.Policy{Type: "cos"}
|
||||
err = fs.DispatchHandler()
|
||||
asserts.NoError(err)
|
||||
|
||||
fs.Policy = &model.Policy{Type: "s3"}
|
||||
err = fs.DispatchHandler()
|
||||
asserts.NoError(err)
|
||||
@@ -262,3 +269,40 @@ func TestFileSystem_SetTargetByInterface(t *testing.T) {
|
||||
asserts.Len(fs.FileTarget, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileSystem_SwitchToSlaveHandler(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
fs := FileSystem{
|
||||
User: &model.User{},
|
||||
}
|
||||
mockNode := &cluster.MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
fs.SwitchToSlaveHandler(mockNode)
|
||||
a.IsType(&slaveinmaster.Driver{}, fs.Handler)
|
||||
}
|
||||
|
||||
func TestFileSystem_SwitchToShadowHandler(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
fs := FileSystem{
|
||||
User: &model.User{},
|
||||
Policy: &model.Policy{},
|
||||
}
|
||||
mockNode := &cluster.MasterNode{
|
||||
Model: &model.Node{},
|
||||
}
|
||||
|
||||
// local to remote
|
||||
{
|
||||
fs.Policy.Type = "local"
|
||||
fs.SwitchToShadowHandler(mockNode, "", "")
|
||||
a.IsType(&masterinslave.Driver{}, fs.Handler)
|
||||
}
|
||||
|
||||
// onedrive
|
||||
{
|
||||
fs.Policy.Type = "onedrive"
|
||||
fs.SwitchToShadowHandler(mockNode, "", "")
|
||||
a.IsType(&masterinslave.Driver{}, fs.Handler)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -708,3 +708,29 @@ func TestHookGiveBackCapacity(t *testing.T) {
|
||||
asserts.EqualValues(7, fs.User.Storage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHookValidateCapacityWithoutIncrease(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
fs := &FileSystem{
|
||||
User: &model.User{
|
||||
Model: gorm.Model{ID: 1},
|
||||
Storage: 10,
|
||||
Group: model.Group{},
|
||||
},
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, local.FileStream{Size: 1})
|
||||
|
||||
// not enough
|
||||
{
|
||||
fs.User.Group.MaxStorage = 10
|
||||
a.Error(HookValidateCapacityWithoutIncrease(ctx, fs))
|
||||
a.EqualValues(10, fs.User.Storage)
|
||||
}
|
||||
|
||||
// enough
|
||||
{
|
||||
fs.User.Group.MaxStorage = 11
|
||||
a.NoError(HookValidateCapacityWithoutIncrease(ctx, fs))
|
||||
a.EqualValues(10, fs.User.Storage)
|
||||
}
|
||||
}
|
||||
|
||||
+49
-3
@@ -4,6 +4,9 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"sync"
|
||||
|
||||
"runtime"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
@@ -35,18 +38,53 @@ func (fs *FileSystem) GetThumb(ctx context.Context, id uint) (*response.ContentR
|
||||
ctx = context.WithValue(ctx, fsctx.ThumbSizeCtx, [2]uint{w, h})
|
||||
ctx = context.WithValue(ctx, fsctx.FileModelCtx, fs.FileTarget[0])
|
||||
res, err := fs.Handler.Thumb(ctx, fs.FileTarget[0].SourceName)
|
||||
if err == nil && conf.SystemConfig.Mode == "master" {
|
||||
res.MaxAge = model.GetIntSetting("preview_timeout", 60)
|
||||
}
|
||||
|
||||
// 本地存储策略出错时重新生成缩略图
|
||||
if err != nil && fs.Policy.Type == "local" {
|
||||
fs.GenerateThumbnail(ctx, &fs.FileTarget[0])
|
||||
res, err = fs.Handler.Thumb(ctx, fs.FileTarget[0].SourceName)
|
||||
}
|
||||
|
||||
if err == nil && conf.SystemConfig.Mode == "master" {
|
||||
res.MaxAge = model.GetIntSetting("preview_timeout", 60)
|
||||
}
|
||||
|
||||
return res, err
|
||||
}
|
||||
|
||||
// thumbPool 要使用的任务池
|
||||
var thumbPool *Pool
|
||||
var once sync.Once
|
||||
|
||||
// Pool 带有最大配额的任务池
|
||||
type Pool struct {
|
||||
// 容量
|
||||
worker chan int
|
||||
}
|
||||
|
||||
// Init 初始化任务池
|
||||
func getThumbWorker() *Pool {
|
||||
once.Do(func() {
|
||||
maxWorker := conf.ThumbConfig.MaxTaskCount
|
||||
if maxWorker <= 0 {
|
||||
maxWorker = runtime.GOMAXPROCS(0)
|
||||
}
|
||||
thumbPool = &Pool{
|
||||
worker: make(chan int, maxWorker),
|
||||
}
|
||||
util.Log().Debug("初始化Thumb任务队列,WorkerNum = %d", maxWorker)
|
||||
})
|
||||
return thumbPool
|
||||
}
|
||||
func (pool *Pool) addWorker() {
|
||||
pool.worker <- 1
|
||||
util.Log().Debug("Thumb任务队列,addWorker")
|
||||
}
|
||||
func (pool *Pool) releaseWorker() {
|
||||
util.Log().Debug("Thumb任务队列,releaseWorker")
|
||||
<-pool.worker
|
||||
}
|
||||
|
||||
// GenerateThumbnail 尝试为本地策略文件生成缩略图并获取图像原始大小
|
||||
// TODO 失败时,如果之前还有图像信息,则清除
|
||||
func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
|
||||
@@ -65,6 +103,8 @@ func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
|
||||
return
|
||||
}
|
||||
defer source.Close()
|
||||
getThumbWorker().addWorker()
|
||||
defer getThumbWorker().releaseWorker()
|
||||
|
||||
image, err := thumb.NewThumbFromFile(source, file.Name)
|
||||
if err != nil {
|
||||
@@ -79,6 +119,12 @@ func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
|
||||
image.GetThumb(fs.GenerateThumbnailSize(w, h))
|
||||
// 保存到文件
|
||||
err = image.Save(util.RelativePath(file.SourceName + conf.ThumbConfig.FileSuffix))
|
||||
image = nil
|
||||
if conf.ThumbConfig.GCAfterGen {
|
||||
util.Log().Debug("GenerateThumbnail runtime.GC")
|
||||
runtime.GC()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
util.Log().Warning("无法保存缩略图:%s", err)
|
||||
return
|
||||
|
||||
@@ -38,3 +38,12 @@ func TestFileSystem_GetThumb(t *testing.T) {
|
||||
asserts.EqualValues(50, res.MaxAge)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileSystem_ThumbWorker(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
|
||||
asserts.NotPanics(func() {
|
||||
getThumbWorker().addWorker()
|
||||
getThumbWorker().releaseWorker()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -356,47 +356,6 @@ func TestFileSystem_Delete(t *testing.T) {
|
||||
}}
|
||||
ctx := context.Background()
|
||||
|
||||
//全部未成功
|
||||
{
|
||||
// 列出要删除的目录
|
||||
mock.ExpectQuery("SELECT(.+)").
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{"id"}).
|
||||
AddRow(1).
|
||||
AddRow(2).
|
||||
AddRow(3),
|
||||
)
|
||||
mock.ExpectQuery("SELECT(.+)").
|
||||
WithArgs(1, 2, 3).
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{"id", "name", "source_name", "policy_id", "size"}).
|
||||
AddRow(4, "1.txt", "1.txt", 2, 1),
|
||||
)
|
||||
// 查询顶级的文件
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "name", "source_name", "policy_id", "size"}).AddRow(1, "1.txt", "1.txt", 603, 2))
|
||||
mock.ExpectQuery("SELECT(.+)files(.+)").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "policy_id", "source_name"}))
|
||||
// 查找软连接
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
// 查询上传策略
|
||||
mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "type"}).AddRow(603, "local"))
|
||||
// 删除文件记录
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("DELETE(.+)files").
|
||||
WillReturnResult(sqlmock.NewResult(0, 3))
|
||||
mock.ExpectCommit()
|
||||
// 删除对应分享
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("UPDATE(.+)shares").
|
||||
WillReturnResult(sqlmock.NewResult(0, 3))
|
||||
mock.ExpectCommit()
|
||||
|
||||
err := fs.Delete(ctx, []uint{1}, []uint{1}, false)
|
||||
asserts.Error(err)
|
||||
asserts.Equal(203, err.(serializer.AppError).Code)
|
||||
asserts.Equal(uint64(3), fs.User.Storage)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
//全部未成功,强制
|
||||
{
|
||||
fs.CleanTargets()
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package controllermock
|
||||
|
||||
import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
type SlaveControllerMock struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) HandleHeartBeat(pingReq *serializer.NodePingReq) (serializer.NodePingResp, error) {
|
||||
args := s.Called(pingReq)
|
||||
return args.Get(0).(serializer.NodePingResp), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetAria2Instance(s2 string) (common.Aria2, error) {
|
||||
args := s.Called(s2)
|
||||
return args.Get(0).(common.Aria2), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) SendNotification(s3 string, s2 string, message mq.Message) error {
|
||||
args := s.Called(s3, s2, message)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) SubmitTask(s3 string, i interface{}, s2 string, f func(interface{})) error {
|
||||
args := s.Called(s3, i, s2, f)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetMasterInfo(s2 string) (*cluster.MasterInfo, error) {
|
||||
args := s.Called(s2)
|
||||
return args.Get(0).(*cluster.MasterInfo), args.Error(1)
|
||||
}
|
||||
|
||||
func (s SlaveControllerMock) GetOneDriveToken(s2 string, u uint) (string, error) {
|
||||
args := s.Called(s2, u)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package mocks
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/request"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"io"
|
||||
)
|
||||
|
||||
type NodePoolMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (n NodePoolMock) BalanceNodeByFeature(feature string, lb balancer.Balancer) (error, cluster.Node) {
|
||||
args := n.Called(feature, lb)
|
||||
return args.Error(0), args.Get(1).(cluster.Node)
|
||||
}
|
||||
|
||||
func (n NodePoolMock) GetNodeByID(id uint) cluster.Node {
|
||||
args := n.Called(id)
|
||||
if res, ok := args.Get(0).(cluster.Node); ok {
|
||||
return res
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (n NodePoolMock) Add(node *model.Node) {
|
||||
n.Called(node)
|
||||
}
|
||||
|
||||
func (n NodePoolMock) Delete(id uint) {
|
||||
n.Called(id)
|
||||
}
|
||||
|
||||
type NodeMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (n NodeMock) Init(node *model.Node) {
|
||||
n.Called(node)
|
||||
}
|
||||
|
||||
func (n NodeMock) IsFeatureEnabled(feature string) bool {
|
||||
args := n.Called(feature)
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n NodeMock) SubscribeStatusChange(callback func(isActive bool, id uint)) {
|
||||
n.Called(callback)
|
||||
}
|
||||
|
||||
func (n NodeMock) Ping(req *serializer.NodePingReq) (*serializer.NodePingResp, error) {
|
||||
args := n.Called(req)
|
||||
return args.Get(0).(*serializer.NodePingResp), args.Error(1)
|
||||
}
|
||||
|
||||
func (n NodeMock) IsActive() bool {
|
||||
args := n.Called()
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n NodeMock) GetAria2Instance() common.Aria2 {
|
||||
args := n.Called()
|
||||
return args.Get(0).(common.Aria2)
|
||||
}
|
||||
|
||||
func (n NodeMock) ID() uint {
|
||||
args := n.Called()
|
||||
return args.Get(0).(uint)
|
||||
}
|
||||
|
||||
func (n NodeMock) Kill() {
|
||||
n.Called()
|
||||
}
|
||||
|
||||
func (n NodeMock) IsMater() bool {
|
||||
args := n.Called()
|
||||
return args.Bool(0)
|
||||
}
|
||||
|
||||
func (n NodeMock) MasterAuthInstance() auth.Auth {
|
||||
args := n.Called()
|
||||
return args.Get(0).(auth.Auth)
|
||||
}
|
||||
|
||||
func (n NodeMock) SlaveAuthInstance() auth.Auth {
|
||||
args := n.Called()
|
||||
return args.Get(0).(auth.Auth)
|
||||
}
|
||||
|
||||
func (n NodeMock) DBModel() *model.Node {
|
||||
args := n.Called()
|
||||
return args.Get(0).(*model.Node)
|
||||
}
|
||||
|
||||
type Aria2Mock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (a Aria2Mock) Init() error {
|
||||
args := a.Called()
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) CreateTask(task *model.Download, options map[string]interface{}) (string, error) {
|
||||
args := a.Called(task, options)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) Status(task *model.Download) (rpc.StatusInfo, error) {
|
||||
args := a.Called(task)
|
||||
return args.Get(0).(rpc.StatusInfo), args.Error(1)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) Cancel(task *model.Download) error {
|
||||
args := a.Called(task)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) Select(task *model.Download, files []int) error {
|
||||
args := a.Called(task, files)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) GetConfig() model.Aria2Option {
|
||||
args := a.Called()
|
||||
return args.Get(0).(model.Aria2Option)
|
||||
}
|
||||
|
||||
func (a Aria2Mock) DeleteTempFile(download *model.Download) error {
|
||||
args := a.Called(download)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
type TaskPoolMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (t TaskPoolMock) Add(num int) {
|
||||
t.Called(num)
|
||||
}
|
||||
|
||||
func (t TaskPoolMock) Submit(job task.Job) {
|
||||
t.Called(job)
|
||||
}
|
||||
|
||||
type RequestMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (r RequestMock) Request(method, target string, body io.Reader, opts ...request.Option) *request.Response {
|
||||
return r.Called(method, target, body, opts).Get(0).(*request.Response)
|
||||
}
|
||||
@@ -97,13 +97,13 @@ func (c HTTPClient) Request(method, target string, body io.Reader, opts ...Optio
|
||||
}
|
||||
|
||||
if options.masterMeta && conf.SystemConfig.Mode == "master" {
|
||||
req.Header.Add("X-Site-Url", model.GetSiteURL().String())
|
||||
req.Header.Add("X-Site-Id", model.GetSettingByName("siteID"))
|
||||
req.Header.Add("X-Cloudreve-Version", conf.BackendVersion)
|
||||
req.Header.Add("X-Cr-Site-Url", model.GetSiteURL().String())
|
||||
req.Header.Add("X-Cr-Site-Id", model.GetSettingByName("siteID"))
|
||||
req.Header.Add("X-Cr-Cloudreve-Version", conf.BackendVersion)
|
||||
}
|
||||
|
||||
if options.slaveNodeID != "" && conf.SystemConfig.Mode == "slave" {
|
||||
req.Header.Add("X-Node-Id", options.slaveNodeID)
|
||||
req.Header.Add("X-Cr-Node-Id", options.slaveNodeID)
|
||||
}
|
||||
|
||||
if options.contentLength != -1 {
|
||||
|
||||
@@ -3,6 +3,7 @@ package request
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cache"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
@@ -54,16 +55,19 @@ func TestWithContext(t *testing.T) {
|
||||
|
||||
func TestHTTPClient_Request(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
client := HTTPClient{}
|
||||
client := NewClient(WithSlaveMeta("test"))
|
||||
|
||||
// 正常
|
||||
{
|
||||
resp := client.Request(
|
||||
"GET",
|
||||
"http://cloudreveisnotexist.com",
|
||||
"POST",
|
||||
"/test",
|
||||
strings.NewReader(""),
|
||||
WithContentLength(0),
|
||||
WithEndpoint("http://cloudreveisnotexist.com"),
|
||||
WithTimeout(time.Duration(1)*time.Microsecond),
|
||||
WithCredential(auth.HMACAuth{SecretKey: []byte("123")}, 10),
|
||||
WithoutHeader([]string{"origin", "origin"}),
|
||||
)
|
||||
asserts.Error(resp.Err)
|
||||
asserts.Nil(resp.Response)
|
||||
@@ -224,3 +228,11 @@ func TestNopRSCloser_SetFirstFakeChunk(t *testing.T) {
|
||||
rsc.SetContentLength(20)
|
||||
asserts.EqualValues(20, rsc.status.Size)
|
||||
}
|
||||
|
||||
func TestBlackHole(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
cache.Set("setting_reset_after_upload_failed", "true", 0)
|
||||
a.NotPanics(func() {
|
||||
BlackHole(strings.NewReader("TestBlackHole"))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -76,9 +76,8 @@ func BuildFinishedListResponse(tasks []model.Download) Response {
|
||||
}
|
||||
|
||||
// BuildDownloadingResponse 构建正在下载的列表响应
|
||||
func BuildDownloadingResponse(tasks []model.Download) Response {
|
||||
func BuildDownloadingResponse(tasks []model.Download, intervals map[uint]int) Response {
|
||||
resp := make([]DownloadListResponse, 0, len(tasks))
|
||||
interval := model.GetIntSetting("aria2_interval", 10)
|
||||
|
||||
for i := 0; i < len(tasks); i++ {
|
||||
fileName := ""
|
||||
@@ -92,6 +91,11 @@ func BuildDownloadingResponse(tasks []model.Download) Response {
|
||||
tasks[i].StatusInfo.Files[i2].Path = path.Base(tasks[i].StatusInfo.Files[i2].Path)
|
||||
}
|
||||
|
||||
interval := 10
|
||||
if actualInterval, ok := intervals[tasks[i].ID]; ok {
|
||||
interval = actualInterval
|
||||
}
|
||||
|
||||
resp = append(resp, DownloadListResponse{
|
||||
UpdateTime: tasks[i].UpdatedAt,
|
||||
UpdateInterval: interval,
|
||||
|
||||
@@ -82,10 +82,12 @@ func TestBuildDownloadingResponse(t *testing.T) {
|
||||
},
|
||||
}
|
||||
tasks[1].StatusInfo.BitTorrent.Info.Name = "name.txt"
|
||||
tasks[1].ID = 1
|
||||
|
||||
res := BuildDownloadingResponse(tasks).Data.([]DownloadListResponse)
|
||||
res := BuildDownloadingResponse(tasks, map[uint]int{1: 5}).Data.([]DownloadListResponse)
|
||||
asserts.Len(res, 2)
|
||||
asserts.Equal("name1.txt", res[1].Name)
|
||||
asserts.Equal(5, res[1].UpdateInterval)
|
||||
asserts.Equal("name.txt", res[0].Name)
|
||||
asserts.Equal("name.txt", res[0].Info.Files[0].Path)
|
||||
asserts.Equal("name1.txt", res[1].Info.Files[0].Path)
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
package serializer
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewError(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
err := NewError(400, "Bad Request", errors.New("error"))
|
||||
a.Error(err)
|
||||
a.EqualValues(400, err.Code)
|
||||
|
||||
err.WithError(errors.New("error2"))
|
||||
a.Equal("error2", err.RawError.Error())
|
||||
a.Equal("Bad Request", err.Error())
|
||||
|
||||
resp := &Response{
|
||||
Code: 400,
|
||||
Msg: "Bad Request",
|
||||
Error: "error",
|
||||
}
|
||||
err = NewErrorFromResponse(resp)
|
||||
a.Error(err)
|
||||
}
|
||||
|
||||
func TestDBErr(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
resp := DBErr("", nil)
|
||||
a.NotEmpty(resp.Msg)
|
||||
|
||||
resp = ParamErr("", nil)
|
||||
a.NotEmpty(resp.Msg)
|
||||
}
|
||||
|
||||
func TestErr(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
err := NewError(400, "Bad Request", errors.New("error"))
|
||||
resp := Err(400, "", err)
|
||||
a.Equal("Bad Request", resp.Msg)
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package serializer
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewResponseWithGobData(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
type args struct {
|
||||
data interface{}
|
||||
}
|
||||
|
||||
res := NewResponseWithGobData(args{})
|
||||
a.Equal(CodeInternalSetting, res.Code)
|
||||
|
||||
res = NewResponseWithGobData("TestNewResponseWithGobData")
|
||||
a.Equal(0, res.Code)
|
||||
a.NotEmpty(res.Data)
|
||||
}
|
||||
|
||||
func TestResponse_GobDecode(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
res := NewResponseWithGobData("TestResponse_GobDecode")
|
||||
jsonContent, err := json.Marshal(res)
|
||||
a.NoError(err)
|
||||
resDecoded := &Response{}
|
||||
a.NoError(json.Unmarshal(jsonContent, resDecoded))
|
||||
var target string
|
||||
resDecoded.GobDecode(&target)
|
||||
a.Equal("TestResponse_GobDecode", target)
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package serializer
|
||||
|
||||
import (
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSlaveTransferReq_Hash(t *testing.T) {
|
||||
a := assert.New(t)
|
||||
s1 := &SlaveTransferReq{
|
||||
Src: "1",
|
||||
Policy: &model.Policy{},
|
||||
}
|
||||
s2 := &SlaveTransferReq{
|
||||
Src: "2",
|
||||
Policy: &model.Policy{},
|
||||
}
|
||||
a.NotEqual(s1.Hash("1"), s2.Hash("1"))
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
package slave
|
||||
|
||||
import "github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
|
||||
var (
|
||||
ErrMasterNotFound = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
|
||||
)
|
||||
@@ -144,7 +144,7 @@ func TestCompressTask_Do(t *testing.T) {
|
||||
task.Do()
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.NotEmpty(task.GetError().Msg)
|
||||
asserts.True(util.IsEmpty("test/compress"))
|
||||
asserts.True(util.IsEmpty(util.RelativePath("test/compress")))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+2
-2
@@ -82,7 +82,7 @@ func Record(job Job) (*model.Task, error) {
|
||||
}
|
||||
|
||||
// Resume 从数据库中恢复未完成任务
|
||||
func Resume() {
|
||||
func Resume(p Pool) {
|
||||
tasks := model.GetTasksByStatus(Queued, Processing)
|
||||
if len(tasks) == 0 {
|
||||
return
|
||||
@@ -97,7 +97,7 @@ func Resume() {
|
||||
}
|
||||
|
||||
if job != nil {
|
||||
TaskPoll.Submit(job)
|
||||
p.Submit(job)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+34
-2
@@ -2,6 +2,7 @@ package task
|
||||
|
||||
import (
|
||||
"errors"
|
||||
testMock "github.com/stretchr/testify/mock"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
@@ -21,15 +22,46 @@ func TestRecord(t *testing.T) {
|
||||
asserts.NoError(err)
|
||||
}
|
||||
|
||||
type taskPoolMock struct {
|
||||
testMock.Mock
|
||||
}
|
||||
|
||||
func (t taskPoolMock) Add(num int) {
|
||||
t.Called(num)
|
||||
}
|
||||
|
||||
func (t taskPoolMock) Submit(job Job) {
|
||||
t.Called(job)
|
||||
}
|
||||
|
||||
func TestResume(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
mockPool := taskPoolMock{}
|
||||
|
||||
// 没有任务
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued).WillReturnRows(sqlmock.NewRows([]string{"type"}))
|
||||
Resume()
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}))
|
||||
Resume(mockPool)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// 有任务, 类型未知
|
||||
{
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(233))
|
||||
Resume(mockPool)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// 有任务
|
||||
{
|
||||
mockPool.On("Submit", testMock.Anything)
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type", "props"}).AddRow(CompressTaskType, "{}"))
|
||||
mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
|
||||
Resume(mockPool)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
mockPool.AssertExpectations(t)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetJobFromModel(t *testing.T) {
|
||||
|
||||
+16
-11
@@ -7,23 +7,28 @@ import (
|
||||
)
|
||||
|
||||
// TaskPoll 要使用的任务池
|
||||
var TaskPoll *Pool
|
||||
var TaskPoll Pool
|
||||
|
||||
// Pool 带有最大配额的任务池
|
||||
type Pool struct {
|
||||
type Pool interface {
|
||||
Add(num int)
|
||||
Submit(job Job)
|
||||
}
|
||||
|
||||
// AsyncPool 带有最大配额的任务池
|
||||
type AsyncPool struct {
|
||||
// 容量
|
||||
idleWorker chan int
|
||||
}
|
||||
|
||||
// Add 增加可用Worker数量
|
||||
func (pool *Pool) Add(num int) {
|
||||
func (pool *AsyncPool) Add(num int) {
|
||||
for i := 0; i < num; i++ {
|
||||
pool.idleWorker <- 1
|
||||
}
|
||||
}
|
||||
|
||||
// ObtainWorker 阻塞直到获取新的Worker
|
||||
func (pool *Pool) ObtainWorker() Worker {
|
||||
func (pool *AsyncPool) obtainWorker() Worker {
|
||||
select {
|
||||
case <-pool.idleWorker:
|
||||
// 有空闲Worker名额时,返回新Worker
|
||||
@@ -32,32 +37,32 @@ func (pool *Pool) ObtainWorker() Worker {
|
||||
}
|
||||
|
||||
// FreeWorker 添加空闲Worker
|
||||
func (pool *Pool) FreeWorker() {
|
||||
func (pool *AsyncPool) freeWorker() {
|
||||
pool.Add(1)
|
||||
}
|
||||
|
||||
// Submit 开始提交任务
|
||||
func (pool *Pool) Submit(job Job) {
|
||||
func (pool *AsyncPool) Submit(job Job) {
|
||||
go func() {
|
||||
util.Log().Debug("等待获取Worker")
|
||||
worker := pool.ObtainWorker()
|
||||
worker := pool.obtainWorker()
|
||||
util.Log().Debug("获取到Worker")
|
||||
worker.Do(job)
|
||||
util.Log().Debug("释放Worker")
|
||||
pool.FreeWorker()
|
||||
pool.freeWorker()
|
||||
}()
|
||||
}
|
||||
|
||||
// Init 初始化任务池
|
||||
func Init() {
|
||||
maxWorker := model.GetIntSetting("max_worker_num", 10)
|
||||
TaskPoll = &Pool{
|
||||
TaskPoll = &AsyncPool{
|
||||
idleWorker: make(chan int, maxWorker),
|
||||
}
|
||||
TaskPoll.Add(maxWorker)
|
||||
util.Log().Info("初始化任务队列,WorkerNum = %d", maxWorker)
|
||||
|
||||
if conf.SystemConfig.Mode == "master" {
|
||||
Resume()
|
||||
Resume(TaskPoll)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,15 +29,15 @@ func TestMain(m *testing.M) {
|
||||
func TestInit(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
cache.Set("setting_max_worker_num", "10", 0)
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(-1))
|
||||
mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(-1))
|
||||
Init()
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.Len(TaskPoll.idleWorker, 10)
|
||||
asserts.Len(TaskPoll.(*AsyncPool).idleWorker, 10)
|
||||
}
|
||||
|
||||
func TestPool_Submit(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
pool := &Pool{
|
||||
pool := &AsyncPool{
|
||||
idleWorker: make(chan int, 1),
|
||||
}
|
||||
pool.Add(1)
|
||||
|
||||
@@ -3,15 +3,14 @@ package slavetask
|
||||
import (
|
||||
"context"
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/slave"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/task"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// TransferTask 文件中转任务
|
||||
@@ -68,8 +67,8 @@ func (job *TransferTask) SetErrorMsg(msg string, err error) {
|
||||
},
|
||||
}
|
||||
|
||||
if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
|
||||
util.Log().Warning("无法发送转存失败通知到从机, ", err)
|
||||
if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
|
||||
util.Log().Warning("无法发送转存失败通知到从机, %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -94,7 +93,7 @@ func (job *TransferTask) Do() {
|
||||
return
|
||||
}
|
||||
|
||||
master, err := slave.DefaultController.GetMasterInfo(job.MasterID)
|
||||
master, err := cluster.DefaultController.GetMasterInfo(job.MasterID)
|
||||
if err != nil {
|
||||
job.SetErrorMsg("找不到主机节点", err)
|
||||
return
|
||||
@@ -131,15 +130,15 @@ func (job *TransferTask) Do() {
|
||||
Content: serializer.SlaveTransferResult{},
|
||||
}
|
||||
|
||||
if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
|
||||
util.Log().Warning("无法发送转存成功通知到从机, ", err)
|
||||
if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
|
||||
util.Log().Warning("无法发送转存成功通知到从机, %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Recycle 回收临时文件
|
||||
func (job *TransferTask) Recycle() {
|
||||
err := os.RemoveAll(filepath.Dir(job.Req.Src))
|
||||
err := os.Remove(job.Req.Src)
|
||||
if err != nil {
|
||||
util.Log().Warning("无法删除中转临时目录[%s], %s", job.Req.Src, err)
|
||||
util.Log().Warning("无法删除中转临时文件[%s], %s", job.Req.Src, err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,7 +138,7 @@ func TestNewTransferTask(t *testing.T) {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("INSERT(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
job, err := NewTransferTask(1, []string{}, "/", "/", false)
|
||||
job, err := NewTransferTask(1, []string{}, "/", "/", false, 0, nil)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.NotNil(job)
|
||||
asserts.NoError(err)
|
||||
@@ -150,7 +150,7 @@ func TestNewTransferTask(t *testing.T) {
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec("INSERT(.+)").WillReturnError(errors.New("error"))
|
||||
mock.ExpectRollback()
|
||||
job, err := NewTransferTask(1, []string{}, "/", "/", false)
|
||||
job, err := NewTransferTask(1, []string{}, "/", "/", false, 0, nil)
|
||||
asserts.NoError(mock.ExpectationsWereMet())
|
||||
asserts.Nil(job)
|
||||
asserts.Error(err)
|
||||
|
||||
+55
-4
@@ -12,9 +12,11 @@ import (
|
||||
"strings"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
|
||||
"github.com/nfnt/resize"
|
||||
//"github.com/nfnt/resize"
|
||||
"golang.org/x/image/draw"
|
||||
)
|
||||
|
||||
// Thumb 缩略图
|
||||
@@ -58,7 +60,8 @@ func NewThumbFromFile(file io.Reader, name string) (*Thumb, error) {
|
||||
|
||||
// GetThumb 生成给定最大尺寸的缩略图
|
||||
func (image *Thumb) GetThumb(width, height uint) {
|
||||
image.src = resize.Thumbnail(width, height, image.src, resize.Lanczos3)
|
||||
//image.src = resize.Thumbnail(width, height, image.src, resize.Lanczos3)
|
||||
image.src = Thumbnail(width, height, image.src)
|
||||
}
|
||||
|
||||
// GetSize 获取图像尺寸
|
||||
@@ -75,12 +78,59 @@ func (image *Thumb) Save(path string) (err error) {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
switch conf.ThumbConfig.EncodeMethod {
|
||||
case "png":
|
||||
err = png.Encode(out, image.src)
|
||||
default:
|
||||
err = jpeg.Encode(out, image.src, &jpeg.Options{Quality: conf.ThumbConfig.EncodeQuality})
|
||||
}
|
||||
|
||||
err = png.Encode(out, image.src)
|
||||
return err
|
||||
|
||||
}
|
||||
|
||||
// Thumbnail will downscale provided image to max width and height preserving
|
||||
// original aspect ratio and using the interpolation function interp.
|
||||
// It will return original image, without processing it, if original sizes
|
||||
// are already smaller than provided constraints.
|
||||
func Thumbnail(maxWidth, maxHeight uint, img image.Image) image.Image {
|
||||
origBounds := img.Bounds()
|
||||
origWidth := uint(origBounds.Dx())
|
||||
origHeight := uint(origBounds.Dy())
|
||||
newWidth, newHeight := origWidth, origHeight
|
||||
|
||||
// Return original image if it have same or smaller size as constraints
|
||||
if maxWidth >= origWidth && maxHeight >= origHeight {
|
||||
return img
|
||||
}
|
||||
|
||||
// Preserve aspect ratio
|
||||
if origWidth > maxWidth {
|
||||
newHeight = uint(origHeight * maxWidth / origWidth)
|
||||
if newHeight < 1 {
|
||||
newHeight = 1
|
||||
}
|
||||
newWidth = maxWidth
|
||||
}
|
||||
|
||||
if newHeight > maxHeight {
|
||||
newWidth = uint(newWidth * maxHeight / newHeight)
|
||||
if newWidth < 1 {
|
||||
newWidth = 1
|
||||
}
|
||||
newHeight = maxHeight
|
||||
}
|
||||
return Resize(newWidth, newHeight, img)
|
||||
}
|
||||
|
||||
func Resize(newWidth, newHeight uint, img image.Image) image.Image {
|
||||
// Set the expected size that you want:
|
||||
dst := image.NewRGBA(image.Rect(0, 0, int(newWidth), int(newHeight)))
|
||||
// Resize:
|
||||
draw.BiLinear.Scale(dst, dst.Rect, img, img.Bounds(), draw.Src, nil)
|
||||
return dst
|
||||
}
|
||||
|
||||
// CreateAvatar 创建头像
|
||||
func (image *Thumb) CreateAvatar(uid uint) error {
|
||||
// 读取头像相关设定
|
||||
@@ -92,7 +142,8 @@ func (image *Thumb) CreateAvatar(uid uint) error {
|
||||
// 生成头像缩略图
|
||||
src := image.src
|
||||
for k, size := range []int{s, m, l} {
|
||||
image.src = resize.Resize(uint(size), uint(size), src, resize.Lanczos3)
|
||||
//image.src = resize.Resize(uint(size), uint(size), src, resize.Lanczos3)
|
||||
image.src = Resize(uint(size), uint(size), src)
|
||||
err := image.Save(filepath.Join(savePath, fmt.Sprintf("avatar_%d_%d.png", uid, k)))
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -86,6 +86,30 @@ func TestThumb_GetThumb(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestThumb_Thumbnail(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
{
|
||||
img := image.NewRGBA(image.Rect(0, 0, 500, 200))
|
||||
thumb := Thumbnail(100, 100, img)
|
||||
asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 40))
|
||||
}
|
||||
{
|
||||
img := image.NewRGBA(image.Rect(0, 0, 200, 200))
|
||||
thumb := Thumbnail(100, 100, img)
|
||||
asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 100))
|
||||
}
|
||||
{
|
||||
img := image.NewRGBA(image.Rect(0, 0, 500, 500))
|
||||
thumb := Thumbnail(100, 100, img)
|
||||
asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 100))
|
||||
}
|
||||
{
|
||||
img := image.NewRGBA(image.Rect(0, 0, 200, 500))
|
||||
thumb := Thumbnail(100, 100, img)
|
||||
asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 40, 100))
|
||||
}
|
||||
}
|
||||
|
||||
func TestThumb_Save(t *testing.T) {
|
||||
asserts := assert.New(t)
|
||||
file := CreateTestImage()
|
||||
|
||||
+20
-17
@@ -38,23 +38,26 @@ func moveFiles(ctx context.Context, fs *filesystem.FileSystem, src FileInfo, dst
|
||||
fileIDs = []uint{src.(*model.File).ID}
|
||||
}
|
||||
|
||||
// 判断是否为重命名
|
||||
if src.GetPosition() == path.Dir(dst) {
|
||||
err = fs.Rename(
|
||||
ctx,
|
||||
folderIDs,
|
||||
fileIDs,
|
||||
path.Base(dst),
|
||||
)
|
||||
} else {
|
||||
err = fs.Move(
|
||||
ctx,
|
||||
folderIDs,
|
||||
fileIDs,
|
||||
src.GetPosition(),
|
||||
path.Dir(dst),
|
||||
)
|
||||
}
|
||||
// 判断是否需要移动
|
||||
if src.GetPosition() != path.Dir(dst) {
|
||||
err = fs.Move(
|
||||
ctx,
|
||||
folderIDs,
|
||||
fileIDs,
|
||||
src.GetPosition(),
|
||||
path.Dir(dst),
|
||||
)
|
||||
}
|
||||
|
||||
// 判断是否需要重命名
|
||||
if err == nil && src.GetName() != path.Base(dst) {
|
||||
err = fs.Rename(
|
||||
ctx,
|
||||
folderIDs,
|
||||
fileIDs,
|
||||
path.Base(dst),
|
||||
)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return http.StatusInternalServerError, err
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package controllers
|
||||
|
||||
import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"io"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
@@ -72,7 +74,7 @@ func AdminReloadService(c *gin.Context) {
|
||||
case "email":
|
||||
email.Init()
|
||||
case "aria2":
|
||||
aria2.Init(true)
|
||||
aria2.Init(true, cluster.Default, mq.GlobalMQ)
|
||||
}
|
||||
|
||||
c.JSON(200, serializer.Response{})
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
package controllers
|
||||
|
||||
import "C"
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
@@ -10,6 +9,7 @@ import (
|
||||
"sync"
|
||||
|
||||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/local"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
|
||||
@@ -180,7 +180,7 @@ func Thumb(c *gin.Context) {
|
||||
}
|
||||
|
||||
defer resp.Content.Close()
|
||||
http.ServeContent(c.Writer, c.Request, "thumb.png", fs.FileTarget[0].UpdatedAt, resp.Content)
|
||||
http.ServeContent(c.Writer, c.Request, "thumb."+conf.ThumbConfig.EncodeMethod, fs.FileTarget[0].UpdatedAt, resp.Content)
|
||||
|
||||
}
|
||||
|
||||
@@ -308,8 +308,8 @@ func FileUploadStream(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解码文件名和路径
|
||||
fileName, err := url.QueryUnescape(c.Request.Header.Get("X-FileName"))
|
||||
filePath, err := url.QueryUnescape(c.Request.Header.Get("X-Path"))
|
||||
fileName, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-FileName"))
|
||||
filePath, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-Path"))
|
||||
if err != nil {
|
||||
c.JSON(200, ErrorResponse(err))
|
||||
return
|
||||
|
||||
@@ -32,7 +32,7 @@ func SlaveUpload(c *gin.Context) {
|
||||
fs.Handler = local.Driver{}
|
||||
|
||||
// 从请求中取得上传策略
|
||||
uploadPolicyRaw := c.GetHeader("X-Policy")
|
||||
uploadPolicyRaw := c.GetHeader("X-Cr-Policy")
|
||||
if uploadPolicyRaw == "" {
|
||||
c.JSON(200, serializer.ParamErr("未指定上传策略", nil))
|
||||
return
|
||||
@@ -54,7 +54,7 @@ func SlaveUpload(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 解码文件名和路径
|
||||
fileName, err := url.QueryUnescape(c.Request.Header.Get("X-FileName"))
|
||||
fileName, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-FileName"))
|
||||
if err != nil {
|
||||
c.JSON(200, ErrorResponse(err))
|
||||
return
|
||||
@@ -74,7 +74,7 @@ func SlaveUpload(c *gin.Context) {
|
||||
fs.Use("AfterValidateFailed", filesystem.HookDeleteTempFile)
|
||||
|
||||
// 是否允许覆盖
|
||||
if c.Request.Header.Get("X-Overwrite") == "false" {
|
||||
if c.Request.Header.Get("X-Cr-Overwrite") == "false" {
|
||||
ctx = context.WithValue(ctx, fsctx.DisableOverwrite, true)
|
||||
}
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ func TestLocalFileUpload(t *testing.T) {
|
||||
strings.NewReader("2333"),
|
||||
)
|
||||
req.Header.Add("Content-Length", "4")
|
||||
req.Header.Add("X-FileName", "大地的%sfsf")
|
||||
req.Header.Add("X-Cr-FileName", "大地的%sfsf")
|
||||
return req
|
||||
},
|
||||
ExpectCode: 40002,
|
||||
@@ -87,8 +87,8 @@ func TestLocalFileUpload(t *testing.T) {
|
||||
strings.NewReader("2333"),
|
||||
)
|
||||
req.Header.Add("Content-Length", "4")
|
||||
req.Header.Add("X-FileName", "TestFileUploadRoute.txt")
|
||||
req.Header.Add("X-Path", "/")
|
||||
req.Header.Add("X-Cr-FileName", "TestFileUploadRoute.txt")
|
||||
req.Header.Add("X-Cr-Path", "/")
|
||||
return req
|
||||
},
|
||||
ExpectCode: 0,
|
||||
|
||||
+3
-2
@@ -3,6 +3,7 @@ package routers
|
||||
import (
|
||||
"github.com/cloudreve/Cloudreve/v3/middleware"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/auth"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/hashid"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
@@ -59,7 +60,7 @@ func InitSlaveRouter() *gin.Engine {
|
||||
|
||||
// 离线下载
|
||||
aria2 := v3.Group("aria2")
|
||||
aria2.Use(middleware.UseSlaveAria2Instance())
|
||||
aria2.Use(middleware.UseSlaveAria2Instance(cluster.DefaultController))
|
||||
{
|
||||
// 创建离线下载任务
|
||||
aria2.POST("task", controllers.SlaveAria2Create)
|
||||
@@ -205,7 +206,7 @@ func InitMasterRouter() *gin.Engine {
|
||||
|
||||
// 从机的 RPC 通信
|
||||
slave := v3.Group("slave")
|
||||
slave.Use(middleware.SlaveRPCSignRequired())
|
||||
slave.Use(middleware.SlaveRPCSignRequired(cluster.Default))
|
||||
{
|
||||
// 事件通知
|
||||
slave.PUT("notification/:subject", controllers.SlaveNotificationPush)
|
||||
|
||||
@@ -24,7 +24,10 @@ func (service *AddNodeService) Add() serializer.Response {
|
||||
}
|
||||
}
|
||||
|
||||
cluster.Default.Add(&service.Node)
|
||||
if service.Node.Status == model.NodeActive {
|
||||
cluster.Default.Add(&service.Node)
|
||||
}
|
||||
|
||||
return serializer.Response{Data: service.Node.ID}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/mq"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/slave"
|
||||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
@@ -49,9 +48,7 @@ func (service *AddURLService) Add(c *gin.Context, taskType int) serializer.Respo
|
||||
}
|
||||
|
||||
// 获取 Aria2 负载均衡器
|
||||
aria2.Lock.RLock()
|
||||
lb := aria2.LB
|
||||
aria2.Lock.RUnlock()
|
||||
lb := aria2.GetLoadBalancer()
|
||||
|
||||
// 获取 Aria2 实例
|
||||
err, node := cluster.Default.BalanceNodeByFeature("aria2", lb)
|
||||
@@ -73,7 +70,7 @@ func (service *AddURLService) Add(c *gin.Context, taskType int) serializer.Respo
|
||||
}
|
||||
|
||||
// 创建任务监控
|
||||
monitor.NewMonitor(task)
|
||||
monitor.NewMonitor(task, cluster.Default, mq.GlobalMQ)
|
||||
|
||||
return serializer.Response{}
|
||||
}
|
||||
@@ -91,7 +88,7 @@ func Add(c *gin.Context, service *serializer.SlaveAria2Call) serializer.Response
|
||||
// 创建事件通知回调
|
||||
siteID, _ := c.Get("MasterSiteID")
|
||||
mq.GlobalMQ.SubscribeCallback(gid, func(message mq.Message) {
|
||||
if err := slave.DefaultController.SendNotification(siteID.(string), message.TriggeredBy, message); err != nil {
|
||||
if err := cluster.DefaultController.SendNotification(siteID.(string), message.TriggeredBy, message); err != nil {
|
||||
util.Log().Warning("无法发送离线下载任务状态变更通知, %s", err)
|
||||
}
|
||||
})
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user