sdk/go/arvadostest/run_servers.go

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