X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/6613ec1e9c705fb5b950611fd160d4a2babed251..aa9507e1633819259794ea4d6cf391dc88621dac:/services/keep-balance/integration_test.go

diff --git a/services/keep-balance/integration_test.go b/services/keep-balance/integration_test.go
index 148b783788..42463a002a 100644
--- a/services/keep-balance/integration_test.go
+++ b/services/keep-balance/integration_test.go
@@ -1,26 +1,36 @@
-package main
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package keepbalance
 
 import (
 	"bytes"
-	"log"
-	"net/http"
+	"context"
+	"io"
 	"os"
 	"strings"
 	"testing"
 	"time"
 
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
-	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
-	"git.curoverse.com/arvados.git/sdk/go/arvadostest"
-	"git.curoverse.com/arvados.git/sdk/go/keepclient"
-
+	"git.arvados.org/arvados.git/lib/config"
+	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/arvadosclient"
+	"git.arvados.org/arvados.git/sdk/go/arvadostest"
+	"git.arvados.org/arvados.git/sdk/go/ctxlog"
+	"git.arvados.org/arvados.git/sdk/go/keepclient"
+	"github.com/jmoiron/sqlx"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/sirupsen/logrus"
 	check "gopkg.in/check.v1"
 )
 
 var _ = check.Suite(&integrationSuite{})
 
 type integrationSuite struct {
-	config     Config
+	config     *arvados.Cluster
+	db         *sqlx.DB
+	client     *arvados.Client
 	keepClient *keepclient.KeepClient
 }
 
@@ -29,17 +39,14 @@ func (s *integrationSuite) SetUpSuite(c *check.C) {
 		c.Skip("-short")
 	}
 	arvadostest.ResetEnv()
-	arvadostest.StartAPI()
 	arvadostest.StartKeep(4, true)
 
 	arv, err := arvadosclient.MakeArvadosClient()
-	arv.ApiToken = arvadostest.DataManagerToken
+	arv.ApiToken = arvadostest.SystemRootToken
+	c.Assert(err, check.IsNil)
+
+	s.keepClient, err = keepclient.MakeKeepClient(arv)
 	c.Assert(err, check.IsNil)
-	s.keepClient = &keepclient.KeepClient{
-		Arvados: arv,
-		Client:  &http.Client{},
-	}
-	c.Assert(s.keepClient.DiscoverKeepServers(), check.IsNil)
 	s.putReplicas(c, "foo", 4)
 	s.putReplicas(c, "bar", 1)
 }
@@ -55,30 +62,43 @@ func (s *integrationSuite) TearDownSuite(c *check.C) {
 		c.Skip("-short")
 	}
 	arvadostest.StopKeep(4)
-	arvadostest.StopAPI()
 }
 
 func (s *integrationSuite) SetUpTest(c *check.C) {
-	s.config = Config{
-		Client: arvados.Client{
-			APIHost:   os.Getenv("ARVADOS_API_HOST"),
-			AuthToken: arvadostest.DataManagerToken,
-			Insecure:  true,
-		},
-		KeepServiceTypes: []string{"disk"},
+	cfg, err := config.NewLoader(nil, ctxlog.TestLogger(c)).Load()
+	c.Assert(err, check.Equals, nil)
+	s.config, err = cfg.GetCluster("")
+	c.Assert(err, check.Equals, nil)
+	s.db, err = sqlx.Open("postgres", s.config.PostgreSQL.Connection.String())
+	c.Assert(err, check.IsNil)
+	s.config.Collections.BalancePeriod = arvados.Duration(time.Second)
+
+	s.client = &arvados.Client{
+		APIHost:   os.Getenv("ARVADOS_API_HOST"),
+		AuthToken: arvadostest.SystemRootToken,
+		Insecure:  true,
 	}
 }
 
 func (s *integrationSuite) TestBalanceAPIFixtures(c *check.C) {
-	var logBuf *bytes.Buffer
+	var logBuf bytes.Buffer
 	for iter := 0; iter < 20; iter++ {
-		logBuf := &bytes.Buffer{}
+		logBuf.Reset()
+		logger := logrus.New()
+		logger.Out = io.MultiWriter(&logBuf, os.Stderr)
 		opts := RunOptions{
-			CommitPulls: true,
-			CommitTrash: true,
-			Logger:      log.New(logBuf, "", log.LstdFlags),
+			CommitPulls:           true,
+			CommitTrash:           true,
+			CommitConfirmedFields: true,
+			Logger:                logger,
+		}
+
+		bal := &Balancer{
+			DB:      s.db,
+			Logger:  logger,
+			Metrics: newMetrics(prometheus.NewRegistry()),
 		}
-		nextOpts, err := (&Balancer{}).Run(s.config, opts)
+		nextOpts, err := bal.Run(context.Background(), s.client, s.config, opts)
 		c.Check(err, check.IsNil)
 		c.Check(nextOpts.SafeRendezvousState, check.Not(check.Equals), "")
 		c.Check(nextOpts.CommitPulls, check.Equals, true)
@@ -91,4 +111,23 @@ func (s *integrationSuite) TestBalanceAPIFixtures(c *check.C) {
 		time.Sleep(200 * time.Millisecond)
 	}
 	c.Check(logBuf.String(), check.Not(check.Matches), `(?ms).*0 replicas (0 blocks, 0 bytes) underreplicated.*`)
+
+	for _, trial := range []struct {
+		uuid    string
+		repl    int
+		classes []string
+	}{
+		{arvadostest.EmptyCollectionUUID, 0, []string{}},
+		{arvadostest.FooCollection, 2, []string{"default"}},                                // "foo" blk
+		{arvadostest.StorageClassesDesiredDefaultConfirmedDefault, 2, []string{"default"}}, // "bar" blk
+		{arvadostest.StorageClassesDesiredArchiveConfirmedDefault, 0, []string{}},          // "bar" blk
+	} {
+		c.Logf("%#v", trial)
+		var coll arvados.Collection
+		s.client.RequestAndDecode(&coll, "GET", "arvados/v1/collections/"+trial.uuid, nil, nil)
+		if c.Check(coll.ReplicationConfirmed, check.NotNil) {
+			c.Check(*coll.ReplicationConfirmed, check.Equals, trial.repl)
+		}
+		c.Check(coll.StorageClassesConfirmed, check.DeepEquals, trial.classes)
+	}
 }