services/keep-balance/main_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "bytes"
   9         "io/ioutil"
  10         "net"
  11         "net/http"
  12         "time"
  13
  14         "git.arvados.org/arvados.git/lib/config"
  15         "git.arvados.org/arvados.git/sdk/go/arvados"
  16         "git.arvados.org/arvados.git/sdk/go/arvadostest"
  17         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  18         "github.com/ghodss/yaml"
  19         check "gopkg.in/check.v1"
  20 )
  21
  22 var _ = check.Suite(&mainSuite{})
  23
  24 type mainSuite struct{}
  25
  26 func (s *mainSuite) TestVersionFlag(c *check.C) {
  27         var stdout, stderr bytes.Buffer
  28         runCommand("keep-balance", []string{"-version"}, nil, &stdout, &stderr)
  29         c.Check(stderr.String(), check.Equals, "")
  30         c.Log(stdout.String())
  31 }
  32
  33 func (s *mainSuite) TestHTTPServer(c *check.C) {
  34         arvadostest.StartKeep(2, true)
  35
  36         ln, err := net.Listen("tcp", ":0")
  37         if err != nil {
  38                 c.Fatal(err)
  39         }
  40         _, p, err := net.SplitHostPort(ln.Addr().String())
  41         c.Check(err, check.IsNil)
  42         ln.Close()
  43         cfg, err := config.NewLoader(nil, ctxlog.TestLogger(c)).Load()
  44         c.Assert(err, check.IsNil)
  45         cluster, err := cfg.GetCluster("")
  46         c.Assert(err, check.IsNil)
  47         cluster.Services.Keepbalance.InternalURLs[arvados.URL{Host: "localhost:" + p, Path: "/"}] = arvados.ServiceInstance{}
  48         cfg.Clusters[cluster.ClusterID] = *cluster
  49         config, err := yaml.Marshal(cfg)
  50         c.Assert(err, check.IsNil)
  51
  52         var stdout bytes.Buffer
  53         go runCommand("keep-balance", []string{"-config", "-"}, bytes.NewBuffer(config), &stdout, &stdout)
  54         done := make(chan struct{})
  55         go func() {
  56                 defer close(done)
  57                 for {
  58                         time.Sleep(time.Second / 10)
  59                         req, err := http.NewRequest(http.MethodGet, "http://:"+p+"/metrics", nil)
  60                         if err != nil {
  61                                 c.Fatal(err)
  62                                 return
  63                         }
  64                         req.Header.Set("Authorization", "Bearer "+cluster.ManagementToken)
  65                         resp, err := http.DefaultClient.Do(req)
  66                         if err != nil {
  67                                 c.Logf("error %s", err)
  68                                 continue
  69                         }
  70                         defer resp.Body.Close()
  71                         if resp.StatusCode != http.StatusOK {
  72                                 c.Logf("http status %d", resp.StatusCode)
  73                                 continue
  74                         }
  75                         buf, err := ioutil.ReadAll(resp.Body)
  76                         if err != nil {
  77                                 c.Logf("read body: %s", err)
  78                                 continue
  79                         }
  80                         c.Check(string(buf), check.Matches, `(?ms).*arvados_keepbalance_sweep_seconds_sum.*`)
  81                         return
  82                 }
  83         }()
  84         select {
  85         case <-done:
  86         case <-time.After(time.Second):
  87                 c.Log(stdout.String())
  88                 c.Fatal("timeout")
  89         }
  90         c.Log(stdout.String())
  91
  92         // Check non-metrics URL that gets passed through to us from
  93         // service.Command
  94         req, err := http.NewRequest(http.MethodGet, "http://:"+p+"/not-metrics", nil)
  95         c.Assert(err, check.IsNil)
  96         resp, err := http.DefaultClient.Do(req)
  97         c.Check(err, check.IsNil)
  98         defer resp.Body.Close()
  99         c.Check(resp.StatusCode, check.Equals, http.StatusNotFound)
 100 }