replicate/internal/testingutil/testingutil.go

package testingutil

import (
	"context"
	"database/sql"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"math/rand/v2"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"testing"

	sftpserver "github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"

	"github.com/hanzoai/replicate"
	"github.com/hanzoai/replicate/file"
	"github.com/hanzoai/replicate/internal"
	"github.com/hanzoai/replicate/nats"
	"github.com/hanzoai/replicate/s3"
	"github.com/hanzoai/replicate/sftp"
	"github.com/hanzoai/replicate/webdav"
)

const (
	defaultTigrisEndpoint = "https://fly.storage.tigris.dev"
	defaultTigrisRegion   = "auto"
	defaultTigrisBucket   = "replicate-dev"
	defaultTigrisPathRoot = "integration-tests"
)

var (
	// Enables integration tests.
	integration = flag.Bool("integration", false, "")
	// Enables specific types of replicas to be tested.
	replicaClientTypes = flag.String("replica-clients", "file", "")
	// Sets the log level for the tests.
	logLevel = flag.String("log.level", "debug", "")
)

// S3 settings
var (
	// Replica client settings
	s3AccessKeyID     = flag.String("s3-access-key-id", os.Getenv("REPLICATE_S3_ACCESS_KEY_ID"), "")
	s3SecretAccessKey = flag.String("s3-secret-access-key", os.Getenv("REPLICATE_S3_SECRET_ACCESS_KEY"), "")
	s3Region          = flag.String("s3-region", os.Getenv("REPLICATE_S3_REGION"), "")
	s3Bucket          = flag.String("s3-bucket", os.Getenv("REPLICATE_S3_BUCKET"), "")
	s3Path            = flag.String("s3-path", os.Getenv("REPLICATE_S3_PATH"), "")
	s3Endpoint        = flag.String("s3-endpoint", os.Getenv("REPLICATE_S3_ENDPOINT"), "")
	s3ForcePathStyle  = flag.Bool("s3-force-path-style", os.Getenv("REPLICATE_S3_FORCE_PATH_STYLE") == "true", "")
	s3SkipVerify      = flag.Bool("s3-skip-verify", os.Getenv("REPLICATE_S3_SKIP_VERIFY") == "true", "")
)

// Tigris settings (S3-compatible)
var (
	tigrisAccessKeyID     = flag.String("tigris-access-key-id", os.Getenv("REPLICATE_TIGRIS_ACCESS_KEY_ID"), "")
	tigrisSecretAccessKey = flag.String("tigris-secret-access-key", os.Getenv("REPLICATE_TIGRIS_SECRET_ACCESS_KEY"), "")
)

// Cloudflare R2 settings (S3-compatible)
var (
	r2AccessKeyID     = flag.String("r2-access-key-id", os.Getenv("REPLICATE_R2_ACCESS_KEY_ID"), "")
	r2SecretAccessKey = flag.String("r2-secret-access-key", os.Getenv("REPLICATE_R2_SECRET_ACCESS_KEY"), "")
	r2Endpoint        = flag.String("r2-endpoint", os.Getenv("REPLICATE_R2_ENDPOINT"), "")
	r2Bucket          = flag.String("r2-bucket", os.Getenv("REPLICATE_R2_BUCKET"), "")
)

// Backblaze B2 settings (S3-compatible)
var (
	b2KeyID          = flag.String("b2-key-id", os.Getenv("REPLICATE_B2_KEY_ID"), "")
	b2ApplicationKey = flag.String("b2-application-key", os.Getenv("REPLICATE_B2_APPLICATION_KEY"), "")
	b2Endpoint       = flag.String("b2-endpoint", os.Getenv("REPLICATE_B2_ENDPOINT"), "")
	b2Bucket         = flag.String("b2-bucket", os.Getenv("REPLICATE_B2_BUCKET"), "")
)

// SFTP settings
var (
	sftpHost     = flag.String("sftp-host", os.Getenv("REPLICATE_SFTP_HOST"), "")
	sftpUser     = flag.String("sftp-user", os.Getenv("REPLICATE_SFTP_USER"), "")
	sftpPassword = flag.String("sftp-password", os.Getenv("REPLICATE_SFTP_PASSWORD"), "")
	sftpKeyPath  = flag.String("sftp-key-path", os.Getenv("REPLICATE_SFTP_KEY_PATH"), "")
	sftpPath     = flag.String("sftp-path", os.Getenv("REPLICATE_SFTP_PATH"), "")
)

// WebDAV settings
var (
	webdavURL      = flag.String("webdav-url", os.Getenv("REPLICATE_WEBDAV_URL"), "")
	webdavUsername = flag.String("webdav-username", os.Getenv("REPLICATE_WEBDAV_USERNAME"), "")
	webdavPassword = flag.String("webdav-password", os.Getenv("REPLICATE_WEBDAV_PASSWORD"), "")
	webdavPath     = flag.String("webdav-path", os.Getenv("REPLICATE_WEBDAV_PATH"), "")
)

// NATS settings
var (
	natsURL      = flag.String("nats-url", os.Getenv("REPLICATE_NATS_URL"), "")
	natsBucket   = flag.String("nats-bucket", os.Getenv("REPLICATE_NATS_BUCKET"), "")
	natsCreds    = flag.String("nats-creds", os.Getenv("REPLICATE_NATS_CREDS"), "")
	natsUsername = flag.String("nats-username", os.Getenv("REPLICATE_NATS_USERNAME"), "")
	natsPassword = flag.String("nats-password", os.Getenv("REPLICATE_NATS_PASSWORD"), "")
)

func Integration() bool {
	return *integration
}

func ReplicaClientTypes() []string {
	return strings.Split(*replicaClientTypes, ",")
}

func NewDB(tb testing.TB, path string) *replicate.DB {
	tb.Helper()
	tb.Logf("db=%s", path)

	level := slog.LevelDebug
	if strings.EqualFold(*logLevel, "trace") {
		level = internal.LevelTrace
	}

	db := replicate.NewDB(path)
	db.Logger = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level:       level,
		ReplaceAttr: internal.ReplaceAttr,
	}))
	return db
}

// MustOpenDBs returns a new instance of a DB & associated SQL DB.
func MustOpenDBs(tb testing.TB) (*replicate.DB, *sql.DB) {
	tb.Helper()
	db := MustOpenDB(tb)
	return db, MustOpenSQLDB(tb, db.Path())
}

// MustCloseDBs closes db & sqldb and removes the parent directory.
func MustCloseDBs(tb testing.TB, db *replicate.DB, sqldb *sql.DB) {
	tb.Helper()
	MustCloseDB(tb, db)
	MustCloseSQLDB(tb, sqldb)
}

// MustOpenDB returns a new instance of a DB.
func MustOpenDB(tb testing.TB) *replicate.DB {
	tb.Helper()
	dir := tb.TempDir()
	return MustOpenDBAt(tb, filepath.Join(dir, "db"))
}

// MustOpenDBAt returns a new instance of a DB for a given path.
func MustOpenDBAt(tb testing.TB, path string) *replicate.DB {
	tb.Helper()
	db := NewDB(tb, path)
	db.MonitorInterval = 0     // disable background goroutine
	db.ShutdownSyncTimeout = 0 // disable shutdown sync retry for faster tests
	db.Replica = replicate.NewReplica(db)
	db.Replica.Client = NewFileReplicaClient(tb)
	db.Replica.MonitorEnabled = false // disable background goroutine
	if err := db.Open(); err != nil {
		tb.Fatal(err)
	}
	return db
}

// MustCloseDB closes db and removes its parent directory.
func MustCloseDB(tb testing.TB, db *replicate.DB) {
	tb.Helper()
	if err := db.Close(context.Background()); err != nil && !strings.Contains(err.Error(), `database is closed`) && !strings.Contains(err.Error(), `file already closed`) {
		tb.Fatal(err)
	} else if err := os.RemoveAll(filepath.Dir(db.Path())); err != nil {
		tb.Fatal(err)
	}
}

// MustOpenSQLDB returns a database/sql DB.
func MustOpenSQLDB(tb testing.TB, path string) *sql.DB {
	tb.Helper()
	d, err := sql.Open("sqlite", path)
	if err != nil {
		tb.Fatal(err)
	} else if _, err := d.ExecContext(context.Background(), `PRAGMA journal_mode = wal;`); err != nil {
		tb.Fatal(err)
	} else if _, err := d.ExecContext(context.Background(), `PRAGMA busy_timeout = 5000;`); err != nil {
		tb.Fatal(err)
	}
	return d
}

// MustCloseSQLDB closes a database/sql DB.
func MustCloseSQLDB(tb testing.TB, d *sql.DB) {
	tb.Helper()
	if err := d.Close(); err != nil {
		tb.Fatal(err)
	}
}

// NewReplicaClient returns a new client for integration testing by type name.
func NewReplicaClient(tb testing.TB, typ string) replicate.ReplicaClient {
	tb.Helper()

	switch typ {
	case file.ReplicaClientType:
		return NewFileReplicaClient(tb)
	case s3.ReplicaClientType:
		return NewS3ReplicaClient(tb)
	case sftp.ReplicaClientType:
		return NewSFTPReplicaClient(tb)
	case webdav.ReplicaClientType:
		return NewWebDAVReplicaClient(tb)
	case nats.ReplicaClientType:
		return NewNATSReplicaClient(tb)
	case "tigris":
		return NewTigrisReplicaClient(tb)
	case "r2":
		return NewR2ReplicaClient(tb)
	case "b2":
		return NewB2ReplicaClient(tb)
	default:
		tb.Fatalf("invalid replica client type: %q", typ)
		return nil
	}
}

// NewFileReplicaClient returns a new client for integration testing.
func NewFileReplicaClient(tb testing.TB) *file.ReplicaClient {
	tb.Helper()
	return file.NewReplicaClient(tb.TempDir())
}

// NewS3ReplicaClient returns a new client for integration testing.
func NewS3ReplicaClient(tb testing.TB) *s3.ReplicaClient {
	tb.Helper()

	c := s3.NewReplicaClient()
	c.AccessKeyID = *s3AccessKeyID
	c.SecretAccessKey = *s3SecretAccessKey
	c.Region = *s3Region
	c.Bucket = *s3Bucket
	c.Path = path.Join(*s3Path, fmt.Sprintf("%016x", rand.Uint64()))
	c.Endpoint = *s3Endpoint
	c.ForcePathStyle = *s3ForcePathStyle
	c.SkipVerify = *s3SkipVerify
	return c
}

// NewTigrisReplicaClient returns an S3 client configured for Fly.io Tigris.
func NewTigrisReplicaClient(tb testing.TB) *s3.ReplicaClient {
	tb.Helper()

	if *tigrisAccessKeyID == "" || *tigrisSecretAccessKey == "" {
		tb.Skip("tigris credentials not configured (set REPLICATE_TIGRIS_ACCESS_KEY_ID/SECRET_ACCESS_KEY)")
	}

	c := s3.NewReplicaClient()
	c.AccessKeyID = *tigrisAccessKeyID
	c.SecretAccessKey = *tigrisSecretAccessKey
	c.Region = defaultTigrisRegion
	c.Bucket = defaultTigrisBucket
	c.Path = path.Join(defaultTigrisPathRoot, fmt.Sprintf("%016x", rand.Uint64()))
	c.Endpoint = defaultTigrisEndpoint
	c.ForcePathStyle = true
	c.RequireContentMD5 = false
	return c
}

// NewR2ReplicaClient returns an S3 client configured for Cloudflare R2.
// Skips the test if R2 credentials are not configured.
func NewR2ReplicaClient(tb testing.TB) *s3.ReplicaClient {
	tb.Helper()

	if *r2AccessKeyID == "" || *r2SecretAccessKey == "" {
		tb.Skip("r2 credentials not configured (set REPLICATE_R2_ACCESS_KEY_ID/SECRET_ACCESS_KEY)")
	}
	if *r2Endpoint == "" {
		tb.Skip("r2 endpoint not configured (set REPLICATE_R2_ENDPOINT)")
	}
	if *r2Bucket == "" {
		tb.Skip("r2 bucket not configured (set REPLICATE_R2_BUCKET)")
	}

	c := s3.NewReplicaClient()
	c.AccessKeyID = *r2AccessKeyID
	c.SecretAccessKey = *r2SecretAccessKey
	c.Region = "auto"
	c.Bucket = *r2Bucket
	c.Path = path.Join("integration-tests", fmt.Sprintf("%016x", rand.Uint64()))
	c.Endpoint = *r2Endpoint
	c.ForcePathStyle = true
	c.SignPayload = true
	return c
}

// NewB2ReplicaClient returns a new Backblaze B2 client for integration testing.
// B2 uses S3-compatible API with path-style URLs and signed payloads.
func NewB2ReplicaClient(tb testing.TB) *s3.ReplicaClient {
	tb.Helper()

	if *b2KeyID == "" || *b2ApplicationKey == "" {
		tb.Skip("b2 credentials not configured (set REPLICATE_B2_KEY_ID/APPLICATION_KEY)")
	}
	if *b2Endpoint == "" {
		tb.Skip("b2 endpoint not configured (set REPLICATE_B2_ENDPOINT)")
	}
	if *b2Bucket == "" {
		tb.Skip("b2 bucket not configured (set REPLICATE_B2_BUCKET)")
	}

	c := s3.NewReplicaClient()
	c.AccessKeyID = *b2KeyID
	c.SecretAccessKey = *b2ApplicationKey
	c.Region = "us-west-002" // B2 uses region in endpoint format
	c.Bucket = *b2Bucket
	c.Path = path.Join("integration-tests", fmt.Sprintf("%016x", rand.Uint64()))
	c.Endpoint = *b2Endpoint
	c.ForcePathStyle = true
	c.SignPayload = true
	return c
}



// NewSFTPReplicaClient returns a new client for integration testing.
func NewSFTPReplicaClient(tb testing.TB) *sftp.ReplicaClient {
	tb.Helper()

	c := sftp.NewReplicaClient()
	c.Host = *sftpHost
	c.User = *sftpUser
	c.Password = *sftpPassword
	c.KeyPath = *sftpKeyPath
	c.Path = path.Join(*sftpPath, fmt.Sprintf("%016x", rand.Uint64()))
	return c
}

// NewWebDAVReplicaClient returns a new client for integration testing.
func NewWebDAVReplicaClient(tb testing.TB) *webdav.ReplicaClient {
	tb.Helper()

	c := webdav.NewReplicaClient()
	c.URL = *webdavURL
	c.Username = *webdavUsername
	c.Password = *webdavPassword
	c.Path = path.Join(*webdavPath, fmt.Sprintf("%016x", rand.Uint64()))
	return c
}

// NewNATSReplicaClient returns a new client for integration testing.
func NewNATSReplicaClient(tb testing.TB) *nats.ReplicaClient {
	tb.Helper()

	c := nats.NewReplicaClient()
	c.URL = *natsURL
	c.BucketName = *natsBucket
	c.Creds = *natsCreds
	c.Username = *natsUsername
	c.Password = *natsPassword
	return c
}


// MustDeleteAll deletes all objects under the client's path.
func MustDeleteAll(tb testing.TB, c replicate.ReplicaClient) {
	tb.Helper()

	if err := c.DeleteAll(context.Background()); err != nil {
		tb.Fatalf("cannot delete all: %s", err)
	}

	switch c := c.(type) {
	case *sftp.ReplicaClient:
		if err := c.Cleanup(context.Background()); err != nil {
			tb.Fatalf("cannot cleanup sftp: %s", err)
		}
	}
}

func MockSFTPServer(t *testing.T, hostKey ssh.Signer) string {
	config := &ssh.ServerConfig{NoClientAuth: true}
	config.AddHostKey(hostKey)

	listener, err := net.Listen("tcp", "127.0.0.1:0") // random available port
	if err != nil {
		t.Fatal(err)
	}

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}

			go func() {
				_, chans, reqs, err := ssh.NewServerConn(conn, config)
				if err != nil {
					return
				}
				go ssh.DiscardRequests(reqs)

				for ch := range chans {
					if ch.ChannelType() != "session" {
						ch.Reject(ssh.UnknownChannelType, "unsupported")
						continue
					}
					channel, requests, err := ch.Accept()
					if err != nil {
						return
					}

					go func(in <-chan *ssh.Request) {
						for req := range in {
							if req.Type == "subsystem" && string(req.Payload[4:]) == "sftp" {
								req.Reply(true, nil)

								server, err := sftpserver.NewServer(channel)
								if err != nil {
									return
								}
								if err := server.Serve(); err != nil && err != io.EOF {
									t.Logf("SFTP server error: %v", err)
								}
								return
							}
							req.Reply(false, nil)
						}
					}(requests)
				}
			}()
		}
	}()

	return listener.Addr().String()
}