provider_test.go 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304
  1. package worker
  2. import (
  3. "fmt"
  4. "io/ioutil"
  5. "os"
  6. "path/filepath"
  7. "testing"
  8. "time"
  9. . "github.com/smartystreets/goconvey/convey"
  10. )
  11. func TestRsyncProvider(t *testing.T) {
  12. Convey("Rsync Provider should work", t, func() {
  13. tmpDir, err := ioutil.TempDir("", "tunasync")
  14. defer os.RemoveAll(tmpDir)
  15. So(err, ShouldBeNil)
  16. scriptFile := filepath.Join(tmpDir, "myrsync")
  17. tmpFile := filepath.Join(tmpDir, "log_file")
  18. c := rsyncConfig{
  19. name: "tuna",
  20. upstreamURL: "rsync://rsync.tuna.moe/tuna/",
  21. rsyncCmd: scriptFile,
  22. workingDir: tmpDir,
  23. logDir: tmpDir,
  24. logFile: tmpFile,
  25. useIPv6: true,
  26. interval: 600 * time.Second,
  27. }
  28. provider, err := newRsyncProvider(c)
  29. So(err, ShouldBeNil)
  30. So(provider.Type(), ShouldEqual, provRsync)
  31. So(provider.Name(), ShouldEqual, c.name)
  32. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  33. So(provider.LogDir(), ShouldEqual, c.logDir)
  34. So(provider.LogFile(), ShouldEqual, c.logFile)
  35. So(provider.Interval(), ShouldEqual, c.interval)
  36. Convey("When entering a context (auto exit)", func() {
  37. func() {
  38. ctx := provider.EnterContext()
  39. defer provider.ExitContext()
  40. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  41. newWorkingDir := "/srv/mirror/working/tuna"
  42. ctx.Set(_WorkingDirKey, newWorkingDir)
  43. So(provider.WorkingDir(), ShouldEqual, newWorkingDir)
  44. }()
  45. Convey("After context is done", func() {
  46. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  47. })
  48. })
  49. Convey("When entering a context (manually exit)", func() {
  50. ctx := provider.EnterContext()
  51. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  52. newWorkingDir := "/srv/mirror/working/tuna"
  53. ctx.Set(_WorkingDirKey, newWorkingDir)
  54. So(provider.WorkingDir(), ShouldEqual, newWorkingDir)
  55. Convey("After context is done", func() {
  56. provider.ExitContext()
  57. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  58. })
  59. })
  60. Convey("Let's try a run", func() {
  61. scriptContent := `#!/bin/bash
  62. echo "syncing to $(pwd)"
  63. echo $@
  64. sleep 1
  65. echo "Done"
  66. exit 0
  67. `
  68. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  69. So(err, ShouldBeNil)
  70. expectedOutput := fmt.Sprintf(
  71. "syncing to %s\n"+
  72. "%s\n"+
  73. "Done\n",
  74. provider.WorkingDir(),
  75. fmt.Sprintf(
  76. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ "+
  77. "--delete --delete-after --delay-updates --safe-links "+
  78. "--timeout=120 --contimeout=120 -6 %s %s",
  79. provider.upstreamURL, provider.WorkingDir(),
  80. ),
  81. )
  82. err = provider.Run()
  83. So(err, ShouldBeNil)
  84. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  85. So(err, ShouldBeNil)
  86. So(string(loggedContent), ShouldEqual, expectedOutput)
  87. // fmt.Println(string(loggedContent))
  88. })
  89. })
  90. }
  91. func TestCmdProvider(t *testing.T) {
  92. Convey("Command Provider should work", t, func(ctx C) {
  93. tmpDir, err := ioutil.TempDir("", "tunasync")
  94. defer os.RemoveAll(tmpDir)
  95. So(err, ShouldBeNil)
  96. scriptFile := filepath.Join(tmpDir, "cmd.sh")
  97. tmpFile := filepath.Join(tmpDir, "log_file")
  98. c := cmdConfig{
  99. name: "tuna-cmd",
  100. upstreamURL: "http://mirrors.tuna.moe/",
  101. command: "bash " + scriptFile,
  102. workingDir: tmpDir,
  103. logDir: tmpDir,
  104. logFile: tmpFile,
  105. interval: 600 * time.Second,
  106. env: map[string]string{
  107. "AOSP_REPO_BIN": "/usr/local/bin/repo",
  108. },
  109. }
  110. provider, err := newCmdProvider(c)
  111. So(err, ShouldBeNil)
  112. So(provider.Type(), ShouldEqual, provCommand)
  113. So(provider.Name(), ShouldEqual, c.name)
  114. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  115. So(provider.LogDir(), ShouldEqual, c.logDir)
  116. So(provider.LogFile(), ShouldEqual, c.logFile)
  117. So(provider.Interval(), ShouldEqual, c.interval)
  118. Convey("Let's try to run a simple command", func() {
  119. scriptContent := `#!/bin/bash
  120. echo $TUNASYNC_WORKING_DIR
  121. echo $TUNASYNC_MIRROR_NAME
  122. echo $TUNASYNC_UPSTREAM_URL
  123. echo $TUNASYNC_LOG_FILE
  124. echo $AOSP_REPO_BIN
  125. `
  126. expectedOutput := fmt.Sprintf(
  127. "%s\n%s\n%s\n%s\n%s\n",
  128. provider.WorkingDir(),
  129. provider.Name(),
  130. provider.upstreamURL,
  131. provider.LogFile(),
  132. "/usr/local/bin/repo",
  133. )
  134. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  135. So(err, ShouldBeNil)
  136. readedScriptContent, err := ioutil.ReadFile(scriptFile)
  137. So(err, ShouldBeNil)
  138. So(readedScriptContent, ShouldResemble, []byte(scriptContent))
  139. err = provider.Run()
  140. So(err, ShouldBeNil)
  141. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  142. So(err, ShouldBeNil)
  143. So(string(loggedContent), ShouldEqual, expectedOutput)
  144. })
  145. Convey("If a command fails", func() {
  146. scriptContent := `exit 1`
  147. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  148. So(err, ShouldBeNil)
  149. readedScriptContent, err := ioutil.ReadFile(scriptFile)
  150. So(err, ShouldBeNil)
  151. So(readedScriptContent, ShouldResemble, []byte(scriptContent))
  152. err = provider.Run()
  153. So(err, ShouldNotBeNil)
  154. })
  155. Convey("If a long job is killed", func(ctx C) {
  156. scriptContent := `#!/bin/bash
  157. sleep 5
  158. `
  159. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  160. So(err, ShouldBeNil)
  161. go func() {
  162. err = provider.Run()
  163. ctx.So(err, ShouldNotBeNil)
  164. }()
  165. time.Sleep(1 * time.Second)
  166. err = provider.Terminate()
  167. So(err, ShouldBeNil)
  168. })
  169. })
  170. }
  171. func TestTwoStageRsyncProvider(t *testing.T) {
  172. Convey("TwoStageRsync Provider should work", t, func(ctx C) {
  173. tmpDir, err := ioutil.TempDir("", "tunasync")
  174. defer os.RemoveAll(tmpDir)
  175. So(err, ShouldBeNil)
  176. scriptFile := filepath.Join(tmpDir, "myrsync")
  177. tmpFile := filepath.Join(tmpDir, "log_file")
  178. c := twoStageRsyncConfig{
  179. name: "tuna-two-stage-rsync",
  180. upstreamURL: "rsync://mirrors.tuna.moe/",
  181. stage1Profile: "debian",
  182. rsyncCmd: scriptFile,
  183. workingDir: tmpDir,
  184. logDir: tmpDir,
  185. logFile: tmpFile,
  186. useIPv6: true,
  187. excludeFile: tmpFile,
  188. }
  189. provider, err := newTwoStageRsyncProvider(c)
  190. So(err, ShouldBeNil)
  191. So(provider.Type(), ShouldEqual, provTwoStageRsync)
  192. So(provider.Name(), ShouldEqual, c.name)
  193. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  194. So(provider.LogDir(), ShouldEqual, c.logDir)
  195. So(provider.LogFile(), ShouldEqual, c.logFile)
  196. So(provider.Interval(), ShouldEqual, c.interval)
  197. Convey("Try a command", func(ctx C) {
  198. scriptContent := `#!/bin/bash
  199. echo "syncing to $(pwd)"
  200. echo $@
  201. sleep 1
  202. echo "Done"
  203. exit 0
  204. `
  205. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  206. So(err, ShouldBeNil)
  207. err = provider.Run()
  208. So(err, ShouldBeNil)
  209. expectedOutput := fmt.Sprintf(
  210. "syncing to %s\n"+
  211. "%s\n"+
  212. "Done\n"+
  213. "syncing to %s\n"+
  214. "%s\n"+
  215. "Done\n",
  216. provider.WorkingDir(),
  217. fmt.Sprintf(
  218. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ --safe-links "+
  219. "--timeout=120 --contimeout=120 --exclude dists/ -6 "+
  220. "--exclude-from %s %s %s",
  221. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  222. ),
  223. provider.WorkingDir(),
  224. fmt.Sprintf(
  225. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ "+
  226. "--delete --delete-after --delay-updates --safe-links "+
  227. "--timeout=120 --contimeout=120 -6 --exclude-from %s %s %s",
  228. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  229. ),
  230. )
  231. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  232. So(err, ShouldBeNil)
  233. So(string(loggedContent), ShouldEqual, expectedOutput)
  234. // fmt.Println(string(loggedContent))
  235. })
  236. Convey("Try terminating", func(ctx C) {
  237. scriptContent := `#!/bin/bash
  238. echo $@
  239. sleep 4
  240. exit 0
  241. `
  242. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  243. So(err, ShouldBeNil)
  244. go func() {
  245. err = provider.Run()
  246. ctx.So(err, ShouldNotBeNil)
  247. }()
  248. time.Sleep(1 * time.Second)
  249. err = provider.Terminate()
  250. So(err, ShouldBeNil)
  251. expectedOutput := fmt.Sprintf(
  252. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ --safe-links "+
  253. "--timeout=120 --contimeout=120 --exclude dists/ -6 "+
  254. "--exclude-from %s %s %s\n",
  255. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  256. )
  257. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  258. So(err, ShouldBeNil)
  259. So(string(loggedContent), ShouldEqual, expectedOutput)
  260. // fmt.Println(string(loggedContent))
  261. })
  262. })
  263. }