tester_test.go

package test_test

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	_ "net/http/pprof"
	"strconv"
	"strings"
	"testing"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig"

	"gfx.cafe/gfx/pggat/lib/auth/credentials"
	"gfx.cafe/gfx/pggat/lib/gat"
	"gfx.cafe/gfx/pggat/lib/gat/gatcaddyfile"
	pool_handler "gfx.cafe/gfx/pggat/lib/gat/handlers/pool"
	"gfx.cafe/gfx/pggat/lib/gat/handlers/rewrite_password"
	"gfx.cafe/gfx/pggat/lib/gat/matchers"
	"gfx.cafe/gfx/pggat/lib/gat/pool"
	"gfx.cafe/gfx/pggat/lib/gat/pool/recipe"
	"gfx.cafe/gfx/pggat/lib/gat/poolers/session"
	"gfx.cafe/gfx/pggat/lib/gat/poolers/transaction"
	"gfx.cafe/gfx/pggat/test"
	"gfx.cafe/gfx/pggat/test/tests"
)

type dialer struct {
	Address  string
	Username string
	Password string
	Database string
}

var nextPort int

func randAddress() string {
	nextPort++
	return "/tmp/.s.PGGAT." + strconv.Itoa(nextPort)
}

func resolveNetwork(address string) string {
	if strings.HasPrefix(address, "/") {
		return "unix"
	} else {
		return "tcp"
	}
}

func randPassword() (string, error) {
	var b [20]byte
	_, err := rand.Read(b[:])
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(b[:]), nil
}

func createServer(parent dialer, poolers map[string]caddy.Module) (server gat.ServerConfig, dialers map[string]dialer, err error) {
	address := randAddress()

	server.Listen = []gat.ListenerConfig{
		{
			Address: address,
		},
	}

	var password string
	password, err = randPassword()
	if err != nil {
		return
	}

	server.Routes = append(
		server.Routes,
		gat.RouteConfig{
			Handle: gatcaddyfile.JSONModuleObject(
				&rewrite_password.Module{
					Password: password,
				},
				gatcaddyfile.Handler,
				"handler",
				nil,
			),
		},
	)

	for name, pooler := range poolers {
		p := pool_handler.Module{
			Config: pool_handler.Config{
				Pooler: gatcaddyfile.JSONModuleObject(
					pooler,
					gatcaddyfile.Pooler,
					"pooler",
					nil,
				),

				ServerAddress: parent.Address,

				ServerUsername: parent.Username,
				ServerPassword: parent.Password,
				ServerDatabase: parent.Database,
			},
		}

		server.Routes = append(server.Routes, gat.RouteConfig{
			Match: gatcaddyfile.JSONModuleObject(
				&matchers.Database{
					Database: name,
				},
				gatcaddyfile.Matcher,
				"matcher",
				nil,
			),
			Handle: gatcaddyfile.JSONModuleObject(
				&p,
				gatcaddyfile.Handler,
				"handler",
				nil,
			),
		})

		if dialers == nil {
			dialers = make(map[string]dialer)
		}
		dialers[name] = dialer{
			Address:  address,
			Username: "pooler",
			Password: password,
			Database: name,
		}
	}

	return
}

func daisyChain(config *gat.Config, control dialer, n int) (dialer, error) {
	for i := 0; i < n; i++ {
		poolConfig := pool.ManagementConfig{}
		var pooler caddy.Module
		if i%2 == 0 {
			pooler = &transaction.Module{
				ManagementConfig: poolConfig,
			}
		} else {
			poolConfig.ServerResetQuery = "DISCARD ALL"
			pooler = &session.Module{
				ManagementConfig: poolConfig,
			}
		}

		server, dialers, err := createServer(control, map[string]caddy.Module{
			"pool": pooler,
		})

		if err != nil {
			return dialer{}, err
		}

		control = dialers["pool"]
		config.Servers = append(config.Servers, server)
	}

	return control, nil
}

func TestTester(t *testing.T) {
	control := recipe.Dialer{
		Network:  "tcp",
		Address:  "localhost:5432",
		Username: "postgres",
		Credentials: credentials.Cleartext{
			Username: "postgres",
			Password: "password",
		},
		Database: "postgres",
	}

	config := gat.Config{}

	parent, err := daisyChain(&config, dialer{
		Address:  "localhost:5432",
		Username: "postgres",
		Password: "password",
		Database: "postgres",
	}, 16)
	if err != nil {
		t.Error(err)
		return
	}

	server, dialers, err := createServer(parent, map[string]caddy.Module{
		"transaction": &transaction.Module{},
		"session": &session.Module{
			ManagementConfig: pool.ManagementConfig{
				ServerResetQuery: "discard all",
			},
		},
	})
	if err != nil {
		t.Error(err)
		return
	}

	config.Servers = append(config.Servers, server)

	transactionDialer := recipe.Dialer{
		Network:  resolveNetwork(dialers["transaction"].Address),
		Address:  dialers["transaction"].Address,
		Username: dialers["transaction"].Username,
		Credentials: credentials.FromString(
			dialers["transaction"].Username,
			dialers["transaction"].Password,
		),
		Database: "transaction",
	}
	sessionDialer := recipe.Dialer{
		Network:  resolveNetwork(dialers["transaction"].Address),
		Address:  dialers["session"].Address,
		Username: dialers["session"].Username,
		Credentials: credentials.FromString(
			dialers["session"].Username,
			dialers["session"].Password,
		),
		Database: "session",
	}

	caddyConfig := caddy.Config{
		AppsRaw: caddy.ModuleMap{
			"pggat": caddyconfig.JSON(config, nil),
		},
	}

	if err = caddy.Run(&caddyConfig); err != nil {
		t.Error(err)
		return
	}

	defer func() {
		_ = caddy.Stop()
	}()

	tester := test.NewTester(test.Config{
		Stress: 8,

		Modes: map[string]recipe.Dialer{
			"control":     control,
			"transaction": transactionDialer,
			"session":     sessionDialer,
		},
	})
	if err = tester.Run(
		tests.SimpleQuery,
		tests.Transaction,
		tests.Sync,
		tests.EQP0,
		tests.EQP1,
		tests.EQP2,
		tests.EQP3,
		tests.EQP4,
		tests.EQP5,
		tests.EQP6,
		tests.EQP7,
		tests.EQP8,
		tests.CopyOut0,
		tests.CopyOut1,
		tests.CopyIn0,
		tests.CopyIn1,
		tests.DiscardAll,
	); err != nil {
		fmt.Print(err.Error())
		t.Fail()
	}
}