provider_test.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372
  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 $RSYNC_PASSWORD $@
  64. sleep 1
  65. echo "Done"
  66. exit 0
  67. `
  68. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  69. So(err, ShouldBeNil)
  70. targetDir, _ := filepath.EvalSymlinks(provider.WorkingDir())
  71. expectedOutput := fmt.Sprintf(
  72. "syncing to %s\n"+
  73. "%s\n"+
  74. "Done\n",
  75. targetDir,
  76. fmt.Sprintf(
  77. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ "+
  78. "--delete --delete-after --delay-updates --safe-links "+
  79. "--timeout=120 --contimeout=120 -6 %s %s",
  80. provider.upstreamURL, provider.WorkingDir(),
  81. ),
  82. )
  83. err = provider.Run()
  84. So(err, ShouldBeNil)
  85. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  86. So(err, ShouldBeNil)
  87. So(string(loggedContent), ShouldEqual, expectedOutput)
  88. // fmt.Println(string(loggedContent))
  89. })
  90. })
  91. }
  92. func TestRsyncProviderWithAuthentication(t *testing.T) {
  93. Convey("Rsync Provider with password should work", t, func() {
  94. tmpDir, err := ioutil.TempDir("", "tunasync")
  95. defer os.RemoveAll(tmpDir)
  96. So(err, ShouldBeNil)
  97. scriptFile := filepath.Join(tmpDir, "myrsync")
  98. tmpFile := filepath.Join(tmpDir, "log_file")
  99. c := rsyncConfig{
  100. name: "tuna",
  101. upstreamURL: "rsync://rsync.tuna.moe/tuna/",
  102. rsyncCmd: scriptFile,
  103. username: "tunasync",
  104. password: "tunasyncpassword",
  105. workingDir: tmpDir,
  106. logDir: tmpDir,
  107. logFile: tmpFile,
  108. useIPv6: true,
  109. interval: 600 * time.Second,
  110. }
  111. provider, err := newRsyncProvider(c)
  112. So(err, ShouldBeNil)
  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 a run", func() {
  119. scriptContent := `#!/bin/bash
  120. echo "syncing to $(pwd)"
  121. echo $USER $RSYNC_PASSWORD $@
  122. sleep 1
  123. echo "Done"
  124. exit 0
  125. `
  126. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  127. So(err, ShouldBeNil)
  128. targetDir, _ := filepath.EvalSymlinks(provider.WorkingDir())
  129. expectedOutput := fmt.Sprintf(
  130. "syncing to %s\n"+
  131. "%s\n"+
  132. "Done\n",
  133. targetDir,
  134. fmt.Sprintf(
  135. "%s %s -aHvh --no-o --no-g --stats --exclude .~tmp~/ "+
  136. "--delete --delete-after --delay-updates --safe-links "+
  137. "--timeout=120 --contimeout=120 -6 %s %s",
  138. provider.username, provider.password, provider.upstreamURL, provider.WorkingDir(),
  139. ),
  140. )
  141. err = provider.Run()
  142. So(err, ShouldBeNil)
  143. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  144. So(err, ShouldBeNil)
  145. So(string(loggedContent), ShouldEqual, expectedOutput)
  146. // fmt.Println(string(loggedContent))
  147. })
  148. })
  149. }
  150. func TestCmdProvider(t *testing.T) {
  151. Convey("Command Provider should work", t, func(ctx C) {
  152. tmpDir, err := ioutil.TempDir("", "tunasync")
  153. defer os.RemoveAll(tmpDir)
  154. So(err, ShouldBeNil)
  155. scriptFile := filepath.Join(tmpDir, "cmd.sh")
  156. tmpFile := filepath.Join(tmpDir, "log_file")
  157. c := cmdConfig{
  158. name: "tuna-cmd",
  159. upstreamURL: "http://mirrors.tuna.moe/",
  160. command: "bash " + scriptFile,
  161. workingDir: tmpDir,
  162. logDir: tmpDir,
  163. logFile: tmpFile,
  164. interval: 600 * time.Second,
  165. env: map[string]string{
  166. "AOSP_REPO_BIN": "/usr/local/bin/repo",
  167. },
  168. }
  169. provider, err := newCmdProvider(c)
  170. So(err, ShouldBeNil)
  171. So(provider.Type(), ShouldEqual, provCommand)
  172. So(provider.Name(), ShouldEqual, c.name)
  173. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  174. So(provider.LogDir(), ShouldEqual, c.logDir)
  175. So(provider.LogFile(), ShouldEqual, c.logFile)
  176. So(provider.Interval(), ShouldEqual, c.interval)
  177. Convey("Let's try to run a simple command", func() {
  178. scriptContent := `#!/bin/bash
  179. echo $TUNASYNC_WORKING_DIR
  180. echo $TUNASYNC_MIRROR_NAME
  181. echo $TUNASYNC_UPSTREAM_URL
  182. echo $TUNASYNC_LOG_FILE
  183. echo $AOSP_REPO_BIN
  184. `
  185. expectedOutput := fmt.Sprintf(
  186. "%s\n%s\n%s\n%s\n%s\n",
  187. provider.WorkingDir(),
  188. provider.Name(),
  189. provider.upstreamURL,
  190. provider.LogFile(),
  191. "/usr/local/bin/repo",
  192. )
  193. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  194. So(err, ShouldBeNil)
  195. readedScriptContent, err := ioutil.ReadFile(scriptFile)
  196. So(err, ShouldBeNil)
  197. So(readedScriptContent, ShouldResemble, []byte(scriptContent))
  198. err = provider.Run()
  199. So(err, ShouldBeNil)
  200. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  201. So(err, ShouldBeNil)
  202. So(string(loggedContent), ShouldEqual, expectedOutput)
  203. })
  204. Convey("If a command fails", func() {
  205. scriptContent := `exit 1`
  206. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  207. So(err, ShouldBeNil)
  208. readedScriptContent, err := ioutil.ReadFile(scriptFile)
  209. So(err, ShouldBeNil)
  210. So(readedScriptContent, ShouldResemble, []byte(scriptContent))
  211. err = provider.Run()
  212. So(err, ShouldNotBeNil)
  213. })
  214. Convey("If a long job is killed", func(ctx C) {
  215. scriptContent := `#!/bin/bash
  216. sleep 5
  217. `
  218. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  219. So(err, ShouldBeNil)
  220. go func() {
  221. err = provider.Run()
  222. ctx.So(err, ShouldNotBeNil)
  223. }()
  224. time.Sleep(1 * time.Second)
  225. err = provider.Terminate()
  226. So(err, ShouldBeNil)
  227. })
  228. })
  229. }
  230. func TestTwoStageRsyncProvider(t *testing.T) {
  231. Convey("TwoStageRsync Provider should work", t, func(ctx C) {
  232. tmpDir, err := ioutil.TempDir("", "tunasync")
  233. defer os.RemoveAll(tmpDir)
  234. So(err, ShouldBeNil)
  235. scriptFile := filepath.Join(tmpDir, "myrsync")
  236. tmpFile := filepath.Join(tmpDir, "log_file")
  237. c := twoStageRsyncConfig{
  238. name: "tuna-two-stage-rsync",
  239. upstreamURL: "rsync://mirrors.tuna.moe/",
  240. stage1Profile: "debian",
  241. rsyncCmd: scriptFile,
  242. workingDir: tmpDir,
  243. logDir: tmpDir,
  244. logFile: tmpFile,
  245. useIPv6: true,
  246. excludeFile: tmpFile,
  247. }
  248. provider, err := newTwoStageRsyncProvider(c)
  249. So(err, ShouldBeNil)
  250. So(provider.Type(), ShouldEqual, provTwoStageRsync)
  251. So(provider.Name(), ShouldEqual, c.name)
  252. So(provider.WorkingDir(), ShouldEqual, c.workingDir)
  253. So(provider.LogDir(), ShouldEqual, c.logDir)
  254. So(provider.LogFile(), ShouldEqual, c.logFile)
  255. So(provider.Interval(), ShouldEqual, c.interval)
  256. Convey("Try a command", func(ctx C) {
  257. scriptContent := `#!/bin/bash
  258. echo "syncing to $(pwd)"
  259. echo $@
  260. sleep 1
  261. echo "Done"
  262. exit 0
  263. `
  264. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  265. So(err, ShouldBeNil)
  266. err = provider.Run()
  267. So(err, ShouldBeNil)
  268. targetDir, _ := filepath.EvalSymlinks(provider.WorkingDir())
  269. expectedOutput := fmt.Sprintf(
  270. "syncing to %s\n"+
  271. "%s\n"+
  272. "Done\n"+
  273. "syncing to %s\n"+
  274. "%s\n"+
  275. "Done\n",
  276. targetDir,
  277. fmt.Sprintf(
  278. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ --safe-links "+
  279. "--timeout=120 --contimeout=120 --exclude dists/ -6 "+
  280. "--exclude-from %s %s %s",
  281. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  282. ),
  283. targetDir,
  284. fmt.Sprintf(
  285. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ "+
  286. "--delete --delete-after --delay-updates --safe-links "+
  287. "--timeout=120 --contimeout=120 -6 --exclude-from %s %s %s",
  288. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  289. ),
  290. )
  291. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  292. So(err, ShouldBeNil)
  293. So(string(loggedContent), ShouldEqual, expectedOutput)
  294. // fmt.Println(string(loggedContent))
  295. })
  296. Convey("Try terminating", func(ctx C) {
  297. scriptContent := `#!/bin/bash
  298. echo $@
  299. sleep 4
  300. exit 0
  301. `
  302. err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
  303. So(err, ShouldBeNil)
  304. go func() {
  305. err = provider.Run()
  306. ctx.So(err, ShouldNotBeNil)
  307. }()
  308. time.Sleep(1 * time.Second)
  309. err = provider.Terminate()
  310. So(err, ShouldBeNil)
  311. expectedOutput := fmt.Sprintf(
  312. "-aHvh --no-o --no-g --stats --exclude .~tmp~/ --safe-links "+
  313. "--timeout=120 --contimeout=120 --exclude dists/ -6 "+
  314. "--exclude-from %s %s %s\n",
  315. provider.excludeFile, provider.upstreamURL, provider.WorkingDir(),
  316. )
  317. loggedContent, err := ioutil.ReadFile(provider.LogFile())
  318. So(err, ShouldBeNil)
  319. So(string(loggedContent), ShouldEqual, expectedOutput)
  320. // fmt.Println(string(loggedContent))
  321. })
  322. })
  323. }