From a5b72b8c5539cdeb78321c12419b5edbe0d5d813 Mon Sep 17 00:00:00 2001
From: Shengqi Chen <harry-chen@outlook.com>
Date: Fri, 28 Feb 2025 14:44:20 +0800
Subject: [PATCH] Add tests for success_exit_codes in config and provider

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
---
 worker/config_test.go   | 56 +++++++++++++++++++++++++++++++++++++++++
 worker/provider_test.go | 53 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/worker/config_test.go b/worker/config_test.go
index 19d85c9..5b564ee 100644
--- a/worker/config_test.go
+++ b/worker/config_test.go
@@ -521,4 +521,60 @@ rsync_options = ["--local"]
 			"--local",       // from mirror.rsync_options
 		})
 	})
+
+	Convey("success_exit_codes should work globally and per mirror", t, func() {
+		tmpfile, err := os.CreateTemp("", "tunasync")
+		So(err, ShouldEqual, nil)
+		defer os.Remove(tmpfile.Name())
+
+		cfgBlob1 := `
+[global]
+name = "test_worker"
+log_dir = "/var/log/tunasync/{{.Name}}"
+mirror_dir = "/data/mirrors"
+concurrent = 10
+interval = 240
+retry = 3
+timeout = 86400
+dangerous_global_success_exit_codes = [10, 20]
+
+[manager]
+api_base = "https://127.0.0.1:5000"
+token = "some_token"
+
+[server]
+hostname = "worker1.example.com"
+listen_addr = "127.0.0.1"
+listen_port = 6000
+ssl_cert = "/etc/tunasync.d/worker1.cert"
+ssl_key = "/etc/tunasync.d/worker1.key"
+
+[[mirrors]]
+name = "foo"
+provider = "rsync"
+upstream = "rsync://foo.bar/"
+interval = 720
+retry = 2
+timeout = 3600
+mirror_dir = "/data/foo"
+success_exit_codes = [30, 40]
+`
+
+		err = os.WriteFile(tmpfile.Name(), []byte(cfgBlob1), 0644)
+		So(err, ShouldEqual, nil)
+		defer tmpfile.Close()
+
+		cfg, err := LoadConfig(tmpfile.Name())
+		So(err, ShouldBeNil)
+
+		providers := map[string]mirrorProvider{}
+		for _, m := range cfg.Mirrors {
+			p := newMirrorProvider(m, cfg)
+			providers[p.Name()] = p
+		}
+
+		p, ok := providers["foo"].(*rsyncProvider)
+		So(ok, ShouldBeTrue)
+		So(p.successExitCodes, ShouldResemble, []int{10, 20, 30, 40})
+	})
 }
diff --git a/worker/provider_test.go b/worker/provider_test.go
index d87bc82..1d6bbfe 100644
--- a/worker/provider_test.go
+++ b/worker/provider_test.go
@@ -552,6 +552,59 @@ sleep 10
 			So(provider.DataSize(), ShouldBeEmpty)
 		})
 	})
+	Convey("Command Provider with successExitCodes should work", t, func(ctx C) {
+		tmpDir, err := os.MkdirTemp("", "tunasync")
+		defer os.RemoveAll(tmpDir)
+		So(err, ShouldBeNil)
+		scriptFile := filepath.Join(tmpDir, "cmd.sh")
+		tmpFile := filepath.Join(tmpDir, "log_file")
+
+		c := cmdConfig{
+			name:        "tuna-cmd",
+			upstreamURL: "http://mirrors.tuna.moe/",
+			command:     "bash " + scriptFile,
+			workingDir:  tmpDir,
+			logDir:      tmpDir,
+			logFile:     tmpFile,
+			interval:    600 * time.Second,
+		}
+
+		provider, err := newCmdProvider(c)
+		provider.SetSuccessExitCodes([]int{199, 200})
+		So(err, ShouldBeNil)
+
+		So(provider.Type(), ShouldEqual, provCommand)
+		So(provider.Name(), ShouldEqual, c.name)
+		So(provider.WorkingDir(), ShouldEqual, c.workingDir)
+		So(provider.LogDir(), ShouldEqual, c.logDir)
+		So(provider.LogFile(), ShouldEqual, c.logFile)
+		So(provider.Interval(), ShouldEqual, c.interval)
+		So(provider.GetSuccessExitCodes(), ShouldResemble, []int{199, 200})
+
+		Convey("Command exits with configured successExitCodes", func() {
+			scriptContent := `exit 199`
+			err = os.WriteFile(scriptFile, []byte(scriptContent), 0755)
+			So(err, ShouldBeNil)
+			readedScriptContent, err := os.ReadFile(scriptFile)
+			So(err, ShouldBeNil)
+			So(readedScriptContent, ShouldResemble, []byte(scriptContent))
+
+			err = provider.Run(make(chan empty, 1))
+			So(err, ShouldBeNil)
+		})
+
+		Convey("Command exits with unknown exit code", func() {
+			scriptContent := `exit 201`
+			err = os.WriteFile(scriptFile, []byte(scriptContent), 0755)
+			So(err, ShouldBeNil)
+			readedScriptContent, err := os.ReadFile(scriptFile)
+			So(err, ShouldBeNil)
+			So(readedScriptContent, ShouldResemble, []byte(scriptContent))
+
+			err = provider.Run(make(chan empty, 1))
+			So(err, ShouldNotBeNil)
+		})
+	})
 }
 
 func TestTwoStageRsyncProvider(t *testing.T) {