featurebase/dax/server/test/managed.go

// Copyright 2021 Molecula Corp. All rights reserved.
package test

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/featurebasedb/featurebase/v3/dax"
	computersvc "github.com/featurebasedb/featurebase/v3/dax/computer/service"
	"github.com/featurebasedb/featurebase/v3/dax/controller"
	controllersvc "github.com/featurebasedb/featurebase/v3/dax/controller/service"
	"github.com/featurebasedb/featurebase/v3/dax/queryer"
	queryersvc "github.com/featurebasedb/featurebase/v3/dax/queryer/service"
	"github.com/featurebasedb/featurebase/v3/dax/server"
	"github.com/featurebasedb/featurebase/v3/errors"
	fbtest "github.com/featurebasedb/featurebase/v3/test"
	"github.com/stretchr/testify/assert"
)

// ManagedCommand represents a test wrapper for server.Command.
type ManagedCommand struct {
	*server.Command

	svcmgr *dax.ServiceManager

	started bool
}

// Manage returns the ServiceManager for the ManagedCommand.
func (mc *ManagedCommand) Manage() *dax.ServiceManager {
	return mc.svcmgr
}

// Address returns the advertise address at which the command can be reached.
func (mc *ManagedCommand) Address() dax.Address {
	uri := mc.URI()
	return dax.Address(uri.String())
}

// Start starts the embedded command.
func (mc *ManagedCommand) Start() error {
	if mc.started {
		return nil
	}

	if err := mc.Command.Start(); err != nil {
		return errors.Wrap(err, "starting command")
	}
	mc.started = true

	return nil
}

// Close closes the embedded  command.
func (mc *ManagedCommand) Close() error {
	return mc.Command.Close()
}

// NewController adds a new ControllerService to the ManagedCommands ServiceManager.
func (mc *ManagedCommand) NewController(cfg controller.Config) dax.ServiceKey {
	uri := mc.URI()
	cfg.Logger = mc.svcmgr.Logger
	mc.svcmgr.Controller = controllersvc.New(uri, cfg)

	return dax.ServicePrefixController
}

// NewQueryer adds a new QueryerService to the ManagedCommands ServiceManager.
func (mc *ManagedCommand) NewQueryer(cfg queryer.Config) dax.ServiceKey {
	uri := mc.URI()
	logger := mc.svcmgr.Logger
	cfg.Logger = logger
	mc.svcmgr.Queryer = queryersvc.New(uri, queryer.New(cfg), logger)

	var controllerAddr dax.Address
	if cfg.ControllerAddress != "" {
		controllerAddr = dax.Address(cfg.ControllerAddress + "/" + dax.ServicePrefixController)
	} else if mc.svcmgr.Controller != nil {
		controllerAddr = mc.svcmgr.Controller.Address()
	}

	// Set Controller
	if err := mc.svcmgr.Queryer.SetController(controllerAddr); err != nil {
		logger.Panicf(errors.Wrap(err, "setting controller").Error())
	}

	return dax.ServicePrefixQueryer
}

// NewComputer adds a new ComputerService to the ManagedCommands ServiceManager.
func (mc *ManagedCommand) NewComputer() dax.ServiceKey {
	cfg := computersvc.CommandConfig{
		ComputerConfig: mc.Config.Computer.Config,

		RootDataDir: mc.Config.Computer.Config.DataDir,

		Stderr: os.Stderr,
		Logger: mc.svcmgr.Logger,
	}

	cfg.ComputerConfig.ControllerAddress = mc.svcmgr.Controller.Address().String()

	// Add new computer service.
	return mc.svcmgr.AddComputer(
		computersvc.New(mc.Address(), cfg, cfg.Logger))
}

// Healthy returns true if the provided service's /health endpoint returns 200
// OK. This means that the service has been added to the ServiceManager and
// started, and that its http handler has been dynamically added.
func (mc *ManagedCommand) Healthy(key dax.ServiceKey) bool {
	if key == "" {
		return false
	}

	addr := mc.Address()
	url := fmt.Sprintf("%s/%s/health", addr.WithScheme("http"), key)
	log.Println("HEALTH URL:", url)
	res, err := http.Get(url)
	if err != nil {
		return false
	} else if res.StatusCode != http.StatusOK {
		return false
	}

	return true
}

// WaitForApplied is a test helper function which retries a computer's
// /directive endpoint a specified number of times, along with a sleep time in
// between tries, until the computer returns applied=true.
func (mc *ManagedCommand) WaitForApplied(t *testing.T, key dax.ServiceKey, n int, sleep time.Duration) {
	t.Helper()

	addr := mc.Address()
	url := fmt.Sprintf("%s/%s/directive", addr.WithScheme("http"), key)
	log.Println("WAIT URL:", url)

	for i := 0; i < n; i++ {
		resp, err := http.Get(url)
		assert.NoError(t, err)
		assert.Equal(t, http.StatusOK, resp.StatusCode)

		body := resp.Body
		defer body.Close()

		got := struct {
			Applied bool `json:"applied"`
		}{}
		assert.NoError(t, json.NewDecoder(body).Decode(&got))

		if got.Applied {
			return
		}

		t.Logf("Wait (%d/%d): url: %s (sleep: %s)\n", i, n, url, sleep.String())
		if i < n-1 {
			time.Sleep(sleep)
		}
	}

	// Getting to here means the directive endpoint never returned successfully,
	// so we need to stop the test.
	t.Fatal("WaitForApplied timed out")
}

// NewManagedCommand returns a new instance of Command.
func NewManagedCommand(tb fbtest.DirCleaner, opts ...server.CommandOption) *ManagedCommand {
	path := tb.TempDir()

	svcmgr := dax.NewServiceManager()
	opts = append(opts, server.OptCommandServiceManager(svcmgr))

	mc := &ManagedCommand{}
	output := io.Discard
	if testing.Verbose() {
		output = os.Stderr
	}
	mc.Command = server.NewCommand(output, opts...)
	mc.svcmgr = svcmgr

	mc.Config.Bind = "http://localhost:0"
	mc.Config.Controller.Config.DataDir = path + "/controller"
	mc.Config.Computer.Config.DataDir = path
	mc.Config.Computer.Config.WriteloggerDir = path + "/wl"
	mc.Config.Controller.Config.WriteloggerDir = path + "/wl"
	mc.Config.Computer.Config.SnapshotterDir = path + "/sn"
	mc.Config.Controller.Config.SnapshotterDir = path + "/sn"

	return mc
}

// DefaultConfig includes a single instance of each service type.
func DefaultConfig() *server.Config {
	cfg := server.NewConfig()
	cfg.Verbose = true
	cfg.Controller.Run = true
	cfg.Controller.Config.RegistrationBatchTimeout = 0
	cfg.Queryer.Run = true
	cfg.Computer.Run = true
	cfg.Computer.N = 1
	return cfg
}

// MustRunManagedCommand starts an in-process set of Services based on the
// provided configuration. If no configuration is provided, it will use the
// DefaultConfig which consists of one instance of each service type.
func MustRunManagedCommand(tb testing.TB, opts ...server.CommandOption) *ManagedCommand {
	// If no opts are passed, use the default configuration which includes a
	// single instance of each service type. This is really just meant to keep
	// test code a bit cleaner when it's not necessary to have a custom service
	// configuration.
	var basic bool
	if len(opts) == 0 {
		opts = []server.CommandOption{
			server.OptCommandConfig(DefaultConfig()),
		}
		basic = true
	}

	mc := NewManagedCommand(tb, opts...)

	if err := mc.Start(); err != nil {
		tb.Fatalf("starting managed command: %v", err)
	}

	if basic {
		assert.True(tb, mc.Healthy(dax.ServicePrefixController))
		assert.True(tb, mc.Healthy(dax.ServicePrefixQueryer))
		assert.True(tb, mc.Healthy(dax.ServicePrefixComputer+"0"))
	}

	return mc
}