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         c.Check(stdout.String(), check.Matches, `keep-balance.*\(go1.*\)\n`)
  32 }
  33
  34 func (s *mainSuite) TestHTTPServer(c *check.C) {
  35         arvadostest.StartKeep(2, true)
  36
  37         ln, err := net.Listen("tcp", ":0")
  38         if err != nil {
  39                 c.Fatal(err)
  40         }
  41         _, p, err := net.SplitHostPort(ln.Addr().String())
  42         c.Check(err, check.IsNil)
  43         ln.Close()
  44         cfg, err := config.NewLoader(nil, ctxlog.TestLogger(c)).Load()
  45         c.Assert(err, check.IsNil)
  46         cluster, err := cfg.GetCluster("")
  47         c.Assert(err, check.IsNil)
  48         cluster.Services.Keepbalance.InternalURLs[arvados.URL{Host: "localhost:" + p, Path: "/"}] = arvados.ServiceInstance{}
  49         cfg.Clusters[cluster.ClusterID] = *cluster
  50         config, err := yaml.Marshal(cfg)
  51         c.Assert(err, check.IsNil)
  52
  53         var stdout bytes.Buffer
  54         go runCommand("keep-balance", []string{"-config", "-"}, bytes.NewBuffer(config), &stdout, &stdout)
  55         done := make(chan struct{})
  56         go func() {
  57                 defer close(done)
  58                 for {
  59                         time.Sleep(time.Second / 10)
  60                         req, err := http.NewRequest(http.MethodGet, "http://:"+p+"/metrics", nil)
  61                         if err != nil {
  62                                 c.Fatal(err)
  63                                 return
  64                         }
  65                         req.Header.Set("Authorization", "Bearer "+cluster.ManagementToken)
  66                         resp, err := http.DefaultClient.Do(req)
  67                         if err != nil {
  68                                 c.Logf("error %s", err)
  69                                 continue
  70                         }
  71                         defer resp.Body.Close()
  72                         if resp.StatusCode != http.StatusOK {
  73                                 c.Logf("http status %d", resp.StatusCode)
  74                                 continue
  75                         }
  76                         buf, err := ioutil.ReadAll(resp.Body)
  77                         if err != nil {
  78                                 c.Logf("read body: %s", err)
  79                                 continue
  80                         }
  81                         c.Check(string(buf), check.Matches, `(?ms).*arvados_keepbalance_sweep_seconds_sum.*`)
  82                         return
  83                 }
  84         }()
  85         select {
  86         case <-done:
  87         case <-time.After(time.Second):
  88                 c.Log(stdout.String())
  89                 c.Fatal("timeout")
  90         }
  91         c.Log(stdout.String())
  92
  93         // Check non-metrics URL that gets passed through to us from
  94         // service.Command
  95         req, err := http.NewRequest(http.MethodGet, "http://:"+p+"/not-metrics", nil)
  96         c.Assert(err, check.IsNil)
  97         resp, err := http.DefaultClient.Do(req)
  98         c.Check(err, check.IsNil)
  99         defer resp.Body.Close()
 100         c.Check(resp.StatusCode, check.Equals, http.StatusNotFound)
 101 }