X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/bf4166939c77771642af846cb5372efc8a78659a..ba1937c21efd8d6392c0479579da1fc06443abdd:/services/keep-balance/integration_test.go

diff --git a/services/keep-balance/integration_test.go b/services/keep-balance/integration_test.go
index e4297bfe6d..2e353c92be 100644
--- a/services/keep-balance/integration_test.go
+++ b/services/keep-balance/integration_test.go
@@ -2,10 +2,11 @@
 //
 // SPDX-License-Identifier: AGPL-3.0
 
-package main
+package keepbalance
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"os"
 	"strings"
@@ -18,6 +19,7 @@ import (
 	"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"
@@ -27,6 +29,7 @@ var _ = check.Suite(&integrationSuite{})
 
 type integrationSuite struct {
 	config     *arvados.Cluster
+	db         *sqlx.DB
 	client     *arvados.Client
 	keepClient *keepclient.KeepClient
 }
@@ -36,7 +39,6 @@ func (s *integrationSuite) SetUpSuite(c *check.C) {
 		c.Skip("-short")
 	}
 	arvadostest.ResetEnv()
-	arvadostest.StartAPI()
 	arvadostest.StartKeep(4, true)
 
 	arv, err := arvadosclient.MakeArvadosClient()
@@ -60,7 +62,6 @@ func (s *integrationSuite) TearDownSuite(c *check.C) {
 		c.Skip("-short")
 	}
 	arvadostest.StopKeep(4)
-	arvadostest.StopAPI()
 }
 
 func (s *integrationSuite) SetUpTest(c *check.C) {
@@ -68,6 +69,8 @@ func (s *integrationSuite) SetUpTest(c *check.C) {
 	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{
@@ -84,19 +87,18 @@ func (s *integrationSuite) TestBalanceAPIFixtures(c *check.C) {
 		logger := logrus.New()
 		logger.Out = io.MultiWriter(&logBuf, os.Stderr)
 		opts := RunOptions{
-			CommitPulls: true,
-			CommitTrash: true,
-			Logger:      logger,
+			CommitConfirmedFields: true,
+			Logger:                logger,
 		}
 
 		bal := &Balancer{
+			DB:      s.db,
 			Logger:  logger,
 			Metrics: newMetrics(prometheus.NewRegistry()),
 		}
-		nextOpts, err := bal.Run(s.client, 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)
 		if iter == 0 {
 			c.Check(logBuf.String(), check.Matches, `(?ms).*ChangeSet{Pulls:1.*`)
 			c.Check(logBuf.String(), check.Not(check.Matches), `(?ms).*ChangeSet{.*Trashes:[^0]}*`)
@@ -112,10 +114,10 @@ func (s *integrationSuite) TestBalanceAPIFixtures(c *check.C) {
 		repl    int
 		classes []string
 	}{
-		{arvadostest.EmptyCollectionUUID, 0, []string{"default"}},
-		{arvadostest.FooCollection, 4, []string{"default"}},                                // "foo" blk
+		{arvadostest.EmptyCollectionUUID, 0, []string{}},
+		{arvadostest.FooCollection, 2, []string{"default"}},                                // "foo" blk
 		{arvadostest.StorageClassesDesiredDefaultConfirmedDefault, 2, []string{"default"}}, // "bar" blk
-		{arvadostest.StorageClassesDesiredArchiveConfirmedDefault, 0, []string{"archive"}}, // "bar" blk
+		{arvadostest.StorageClassesDesiredArchiveConfirmedDefault, 0, []string{}},          // "bar" blk
 	} {
 		c.Logf("%#v", trial)
 		var coll arvados.Collection