testplanet.Run (#1222)

docker-compose.yaml

@@ -4,9 +4,12 @@ services:
     image: redis
   test-postgres:
     image: postgres
+    ports:
+    - 5432:5432
     environment:
     - POSTGRES_USER=storj
     - POSTGRES_PASSWORD=storj-pass
+    - POSTGRES_DB=teststorj
 
   satellite:
     image: storjlabs/satellite:${VERSION:-latest}

internal/testplanet/planet.go

@@ -73,9 +73,14 @@ type Config struct {
 
 // Reconfigure allows to change node configurations
 type Reconfigure struct {
-	Bootstrap   func(planet *Planet, index int, config *bootstrap.Config)
+	NewBootstrapDB func(index int) (bootstrap.DB, error)
-	Satellite   func(planet *Planet, index int, config *satellite.Config)
+	Bootstrap      func(index int, config *bootstrap.Config)
-	StorageNode func(planet *Planet, index int, config *storagenode.Config)
+
+	NewSatelliteDB func(index int) (satellite.DB, error)
+	Satellite      func(index int, config *satellite.Config)
+
+	NewStorageNodeDB func(index int) (storagenode.DB, error)
+	StorageNode      func(index int, config *storagenode.Config)
 }
 
 // Planet is a full storj system setup.
@@ -299,7 +304,12 @@ func (planet *Planet) newSatellites(count int) ([]*satellite.Peer, error) {
 			return nil, err
 		}
 
-		db, err := satellitedb.NewInMemory()
+		var db satellite.DB
+		if planet.config.Reconfigure.NewSatelliteDB != nil {
+			db, err = planet.config.Reconfigure.NewSatelliteDB(i)
+		} else {
+			db, err = satellitedb.NewInMemory()
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -378,7 +388,7 @@ func (planet *Planet) newSatellites(count int) ([]*satellite.Peer, error) {
 			},
 		}
 		if planet.config.Reconfigure.Satellite != nil {
-			planet.config.Reconfigure.Satellite(planet, i, &config)
+			planet.config.Reconfigure.Satellite(i, &config)
 		}
 
 		// TODO: for development only
@@ -418,7 +428,12 @@ func (planet *Planet) newStorageNodes(count int) ([]*storagenode.Peer, error) {
 			return nil, err
 		}
 
-		db, err := storagenodedb.NewInMemory(storageDir)
+		var db storagenode.DB
+		if planet.config.Reconfigure.NewStorageNodeDB != nil {
+			db, err = planet.config.Reconfigure.NewStorageNodeDB(i)
+		} else {
+			db, err = storagenodedb.NewInMemory(storageDir)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -459,7 +474,7 @@ func (planet *Planet) newStorageNodes(count int) ([]*storagenode.Peer, error) {
 			},
 		}
 		if planet.config.Reconfigure.StorageNode != nil {
-			planet.config.Reconfigure.StorageNode(planet, i, &config)
+			planet.config.Reconfigure.StorageNode(i, &config)
 		}
 
 		peer, err := storagenode.New(log, identity, db, config)
@@ -492,9 +507,11 @@ func (planet *Planet) newBootstrap() (peer *bootstrap.Peer, err error) {
 		return nil, err
 	}
 
-	db, err := bootstrapdb.NewInMemory(dbDir)
+	var db bootstrap.DB
-	if err != nil {
+	if planet.config.Reconfigure.NewBootstrapDB != nil {
-		return nil, err
+		db, err = planet.config.Reconfigure.NewBootstrapDB(0)
+	} else {
+		db, err = bootstrapdb.NewInMemory(dbDir)
 	}
 
 	err = db.CreateTables()
@@ -524,7 +541,7 @@ func (planet *Planet) newBootstrap() (peer *bootstrap.Peer, err error) {
 		},
 	}
 	if planet.config.Reconfigure.Bootstrap != nil {
-		planet.config.Reconfigure.Bootstrap(planet, 0, &config)
+		planet.config.Reconfigure.Bootstrap(0, &config)
 	}
 
 	peer, err = bootstrap.New(log, identity, db, config)

internal/testplanet/run.go

@@ -0,0 +1,86 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information
+
+package testplanet
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/zeebo/errs"
+	"go.uber.org/zap/zaptest"
+
+	"storj.io/storj/internal/testcontext"
+	"storj.io/storj/satellite"
+	"storj.io/storj/satellite/satellitedb"
+	"storj.io/storj/satellite/satellitedb/satellitedbtest"
+)
+
+// Run runs testplanet in multiple configurations.
+func Run(t *testing.T, config Config, test func(t *testing.T, ctx *testcontext.Context, planet *Planet)) {
+	schemaSuffix := randomSchemaSuffix()
+	t.Log("schema-suffix ", schemaSuffix)
+
+	for _, satelliteDB := range satellitedbtest.Databases() {
+		t.Run(satelliteDB.Name, func(t *testing.T) {
+			ctx := testcontext.New(t)
+			defer ctx.Cleanup()
+
+			if satelliteDB.URL == "" {
+				t.Skipf("Database %s connection string not provided. %s", satelliteDB.Name, satelliteDB.Message)
+			}
+
+			planetConfig := config
+			planetConfig.Reconfigure.NewBootstrapDB = nil
+			planetConfig.Reconfigure.NewSatelliteDB = func(index int) (satellite.DB, error) {
+				schema := strings.ToLower(t.Name() + "-satellite/" + strconv.Itoa(index) + "-" + schemaSuffix)
+				db, err := satellitedb.New(satellitedbtest.WithSchema(satelliteDB.URL, schema))
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				err = db.CreateSchema(schema)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				return &satelliteSchema{
+					DB:     db,
+					schema: schema,
+				}, nil
+			}
+			planetConfig.Reconfigure.NewStorageNodeDB = nil
+
+			planet, err := NewCustom(zaptest.NewLogger(t), planetConfig)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer ctx.Check(planet.Shutdown)
+
+			planet.Start(ctx)
+			test(t, ctx, planet)
+		})
+	}
+}
+
+// satelliteSchema closes database and drops the associated schema
+type satelliteSchema struct {
+	satellite.DB
+	schema string
+}
+
+func (db *satelliteSchema) Close() error {
+	return errs.Combine(
+		db.DB.DropSchema(db.schema),
+		db.DB.Close(),
+	)
+}
+
+func randomSchemaSuffix() string {
+	var data [8]byte
+	_, _ = rand.Read(data[:])
+	return hex.EncodeToString(data[:])
+}

internal/testplanet/run_test.go

@@ -0,0 +1,21 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information
+
+package testplanet_test
+
+import (
+	"testing"
+
+	"storj.io/storj/internal/testcontext"
+	"storj.io/storj/internal/testplanet"
+)
+
+func TestRun(t *testing.T) {
+	testplanet.Run(t, testplanet.Config{
+		SatelliteCount:   1,
+		StorageNodeCount: 1,
+		UplinkCount:      1,
+	}, func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) {
+		t.Log("running test")
+	})
+}

pkg/kademlia/merge_test.go

@@ -36,7 +36,7 @@ func TestMergePlanets(t *testing.T) {
 		StorageNodeCount: 5,
 		Identities:       alpha.Identities(), // avoid using the same pregenerated identities
 		Reconfigure: testplanet.Reconfigure{
-			Bootstrap: func(planet *testplanet.Planet, index int, config *bootstrap.Config) {
+			Bootstrap: func(index int, config *bootstrap.Config) {
 				config.Kademlia.BootstrapAddr = alpha.Bootstrap.Addr()
 
 			},

satellite/peer.go

@@ -52,8 +52,8 @@ type DB interface {
 	// Close closes the database
 	Close() error
 
-	// SetSchema sets the schema
+	// CreateSchema sets the schema
-	SetSchema(schema string) error
+	CreateSchema(schema string) error
 	// DropSchema drops the schema
 	DropSchema(schema string) error
 

satellite/satellitedb/database.go

@@ -4,6 +4,8 @@
 package satellitedb
 
 import (
+	"strconv"
+
 	"github.com/zeebo/errs"
 
 	"storj.io/storj/internal/migrate"
@@ -57,12 +59,11 @@ func NewInMemory() (satellite.DB, error) {
 	return New("sqlite3://file::memory:?mode=memory")
 }
 
-// SetSchema sets the schema
+// CreateSchema creates a schema if it doesn't exist.
-func (db *DB) SetSchema(schema string) error {
+func (db *DB) CreateSchema(schema string) error {
 	switch db.driver {
 	case "postgres":
-		// TODO: proper escaping
+		_, err := db.db.Exec(`create schema if not exists ` + quoteSchema(schema) + `;`)
-		_, err := db.db.Exec("create schema " + schema + "; set search_path to " + schema + ";")
 		return err
 	}
 	return nil
@@ -72,12 +73,17 @@ func (db *DB) SetSchema(schema string) error {
 func (db *DB) DropSchema(schema string) error {
 	switch db.driver {
 	case "postgres":
-		_, err := db.db.Exec("drop schema " + schema + " cascade;")
+		_, err := db.db.Exec(`drop schema ` + quoteSchema(schema) + ` cascade;`)
 		return err
 	}
 	return nil
 }
 
+// quoteSchema quotes schema name such that it can be used in a postgres query
+func quoteSchema(schema string) string {
+	return strconv.QuoteToASCII(schema)
+}
+
 // BandwidthAgreement is a getter for bandwidth agreement repository
 func (db *DB) BandwidthAgreement() bwagreement.DB {
 	return &bandwidthagreement{db: db.db}

satellite/satellitedb/locked.go

@@ -405,6 +405,13 @@ func (m *lockedUsers) Update(ctx context.Context, user *console.User) error {
 	return m.db.Update(ctx, user)
 }
 
+// CreateSchema sets the schema
+func (m *locked) CreateSchema(schema string) error {
+	m.Lock()
+	defer m.Unlock()
+	return m.db.CreateSchema(schema)
+}
+
 // CreateTables initializes the database
 func (m *locked) CreateTables() error {
 	m.Lock()
@@ -563,13 +570,6 @@ func (m *lockedRepairQueue) Peekqueue(ctx context.Context, limit int) ([]pb.Inju
 	return m.db.Peekqueue(ctx, limit)
 }
 
-// SetSchema sets the schema
-func (m *locked) SetSchema(schema string) error {
-	m.Lock()
-	defer m.Unlock()
-	return m.db.SetSchema(schema)
-}
-
 // StatDB returns database for storing node statistics
 func (m *locked) StatDB() statdb.DB {
 	m.Lock()

satellite/satellitedb/satellitedbtest/run.go

@@ -9,7 +9,9 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"flag"
+	"net/url"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/zeebo/errs"
@@ -19,45 +21,65 @@ import (
 )
 
 const (
-	// postgres connstring that works with docker-compose
+	// DefaultPostgresConn is a connstring that works with docker-compose
-	defaultPostgresConn = "postgres://storj:storj-pass@test-postgres/teststorj?sslmode=disable"
+	DefaultPostgresConn = "postgres://storj:storj-pass@test-postgres/teststorj?sslmode=disable"
-	defaultSqliteConn   = "sqlite3://file::memory:?mode=memory"
+	// DefaultSqliteConn is a connstring that is inmemory
+	DefaultSqliteConn = "sqlite3://file::memory:?mode=memory"
 )
 
 var (
-	testPostgres = flag.String("postgres-test-db", os.Getenv("STORJ_POSTGRES_TEST"), "PostgreSQL test database connection string")
+	// TestPostgres is flag for the postgres test database
+	TestPostgres = flag.String("postgres-test-db", os.Getenv("STORJ_POSTGRES_TEST"), "PostgreSQL test database connection string")
 )
 
+// Database describes a test database
+type Database struct {
+	Name    string
+	URL     string
+	Message string
+}
+
+// Databases returns default databases.
+func Databases() []Database {
+	return []Database{
+		{"Sqlite", DefaultSqliteConn, ""},
+		{"Postgres", *TestPostgres, "Postgres flag missing, example: -postgres-test-db=" + DefaultPostgresConn},
+	}
+}
+
+// WithSchema adds schema param to connection string.
+func WithSchema(connstring string, schema string) string {
+	if strings.HasPrefix(connstring, "postgres") {
+		return connstring + "&search_path=" + url.QueryEscape(schema)
+	}
+	return connstring
+}
+
 // Run method will iterate over all supported databases. Will establish
 // connection and will create tables for each DB.
 func Run(t *testing.T, test func(t *testing.T, db satellite.DB)) {
-	for _, dbInfo := range []struct {
+	schemaSuffix := randomSchemaSuffix()
-		dbName    string
+	t.Log("schema-suffix ", schemaSuffix)
-		dbURL     string
+
-		dbMessage string
+	for _, dbInfo := range Databases() {
-	}{
+		t.Run(dbInfo.Name, func(t *testing.T) {
-		{"Sqlite", defaultSqliteConn, ""},
+			if dbInfo.URL == "" {
-		{"Postgres", *testPostgres, "Postgres flag missing, example: -postgres-test-db=" + defaultPostgresConn},
+				t.Skipf("Database %s connection string not provided. %s", dbInfo.Name, dbInfo.Message)
-	} {
-		t.Run(dbInfo.dbName, func(t *testing.T) {
-			if dbInfo.dbURL == "" {
-				t.Skipf("Database %s connection string not provided. %s", dbInfo.dbName, dbInfo.dbMessage)
 			}
 
-			db, err := satellitedb.New(dbInfo.dbURL)
+			schema := strings.ToLower(t.Name() + "-satellite/x-" + schemaSuffix)
+			db, err := satellitedb.New(WithSchema(dbInfo.URL, schema))
 			if err != nil {
 				t.Fatal(err)
 			}
 
-			schemaName := randomSchemaName() // TODO: create schema name based on t.Name()
+			err = db.CreateSchema(schema)
-
-			err = db.SetSchema(schemaName)
 			if err != nil {
 				t.Fatal(err)
 			}
 
 			defer func() {
-				dropErr := db.DropSchema(schemaName)
+				dropErr := db.DropSchema(schema)
 				err := errs.Combine(dropErr, db.Close())
 				if err != nil {
 					t.Fatal(err)
@@ -74,8 +96,8 @@ func Run(t *testing.T, test func(t *testing.T, db satellite.DB)) {
 	}
 }
 
-func randomSchemaName() string {
+func randomSchemaSuffix() string {
 	var data [8]byte
 	_, _ = rand.Read(data[:])
-	return "s" + hex.EncodeToString(data[:])
+	return hex.EncodeToString(data[:])
 }