sdk/go/arvadostest/run_servers.go

   1 package arvadostest
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "log"
   7         "os"
   8         "os/exec"
   9         "strconv"
  10         "strings"
  11 )
  12
  13 var authSettings = make(map[string]string)
  14
  15 func ResetEnv() {
  16         for k, v := range authSettings {
  17                 os.Setenv(k, v)
  18         }
  19 }
  20
  21 func ParseAuthSettings(authScript []byte) {
  22         scanner := bufio.NewScanner(bytes.NewReader(authScript))
  23         for scanner.Scan() {
  24                 line := scanner.Text()
  25                 if 0 != strings.Index(line, "export ") {
  26                         log.Printf("Ignoring: %v", line)
  27                         continue
  28                 }
  29                 toks := strings.SplitN(strings.Replace(line, "export ", "", 1), "=", 2)
  30                 if len(toks) == 2 {
  31                         authSettings[toks[0]] = toks[1]
  32                 } else {
  33                         log.Fatalf("Could not parse: %v", line)
  34                 }
  35         }
  36         log.Printf("authSettings: %v", authSettings)
  37 }
  38
  39 var pythonTestDir string = ""
  40
  41 func chdirToPythonTests() {
  42         if pythonTestDir != "" {
  43                 if err := os.Chdir(pythonTestDir); err != nil {
  44                         log.Fatalf("chdir %s: %s", pythonTestDir, err)
  45                 }
  46                 return
  47         }
  48         for {
  49                 if err := os.Chdir("sdk/python/tests"); err == nil {
  50                         pythonTestDir, err = os.Getwd()
  51                         return
  52                 }
  53                 if parent, err := os.Getwd(); err != nil || parent == "/" {
  54                         log.Fatalf("sdk/python/tests/ not found in any ancestor")
  55                 }
  56                 if err := os.Chdir(".."); err != nil {
  57                         log.Fatal(err)
  58                 }
  59         }
  60 }
  61
  62 func StartAPI() {
  63         cwd, _ := os.Getwd()
  64         defer os.Chdir(cwd)
  65         chdirToPythonTests()
  66
  67         cmd := exec.Command("python", "run_test_server.py", "start", "--auth", "admin")
  68         cmd.Stdin = nil
  69         cmd.Stderr = os.Stderr
  70
  71         authScript, err := cmd.Output()
  72         if err != nil {
  73                 log.Fatalf("%+v: %s", cmd.Args, err)
  74         }
  75         ParseAuthSettings(authScript)
  76         ResetEnv()
  77 }
  78
  79 func StopAPI() {
  80         cwd, _ := os.Getwd()
  81         defer os.Chdir(cwd)
  82         chdirToPythonTests()
  83
  84         bgRun(exec.Command("python", "run_test_server.py", "stop"))
  85 }
  86
  87 // StartKeep starts the given number of keep servers,
  88 // optionally with -enforce-permissions enabled.
  89 // Use numKeepServers = 2 and enforcePermissions = false under all normal circumstances.
  90 func StartKeep(numKeepServers int, enforcePermissions bool) {
  91         cwd, _ := os.Getwd()
  92         defer os.Chdir(cwd)
  93         chdirToPythonTests()
  94
  95         cmdArgs := []string{"run_test_server.py", "start_keep", "--num-keep-servers", strconv.Itoa(numKeepServers)}
  96         if enforcePermissions {
  97                 cmdArgs = append(cmdArgs, "--keep-enforce-permissions")
  98         }
  99
 100         bgRun(exec.Command("python", cmdArgs...))
 101 }
 102
 103 // StopKeep stops keep servers that were started with StartKeep.
 104 // numkeepServers should be the same value that was passed to StartKeep,
 105 // which is 2 under all normal circumstances.
 106 func StopKeep(numKeepServers int) {
 107         cwd, _ := os.Getwd()
 108         defer os.Chdir(cwd)
 109         chdirToPythonTests()
 110
 111         cmd := exec.Command("python", "run_test_server.py", "stop_keep", "--num-keep-servers", strconv.Itoa(numKeepServers))
 112         cmd.Stdin = nil
 113         cmd.Stderr = os.Stderr
 114         cmd.Stdout = os.Stderr
 115         if err := cmd.Run(); err != nil {
 116                 log.Fatalf("%+v: %s", cmd.Args, err)
 117         }
 118 }
 119
 120 // Start cmd, with stderr and stdout redirected to our own
 121 // stderr. Return when the process exits, but do not wait for its
 122 // stderr and stdout to close: any grandchild processes will continue
 123 // writing to our stderr.
 124 func bgRun(cmd *exec.Cmd) {
 125         cmd.Stdin = nil
 126         cmd.Stderr = os.Stderr
 127         cmd.Stdout = os.Stderr
 128         if err := cmd.Start(); err != nil {
 129                 log.Fatalf("%+v: %s", cmd.Args, err)
 130         }
 131         if _, err := cmd.Process.Wait(); err != nil {
 132                 log.Fatalf("%+v: %s", cmd.Args, err)
 133         }
 134 }