setup.go

Documentation: github.com/cybertec-postgresql/pgwatch/v6/internal/testutil

     1  package testutil
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"net"
     7  	"os"
     8  	"time"
     9  
    10  	"github.com/cybertec-postgresql/pgwatch/v6/api/pb"
    11  	"github.com/testcontainers/testcontainers-go"
    12  	"github.com/testcontainers/testcontainers-go/modules/etcd"
    13  	"github.com/testcontainers/testcontainers-go/modules/postgres"
    14  	"github.com/testcontainers/testcontainers-go/wait"
    15  	"google.golang.org/grpc"
    16  	"google.golang.org/grpc/codes"
    17  	"google.golang.org/grpc/credentials"
    18  	"google.golang.org/grpc/metadata"
    19  	"google.golang.org/grpc/status"
    20  )
    21  
    22  func SetupPostgresContainer() (*postgres.PostgresContainer, func(), error) {
    23  	pgContainer, err := postgres.Run(TestContext,
    24  		PostgresImage,
    25  		postgres.WithDatabase(MockDatabase),
    26  		testcontainers.WithWaitStrategy(
    27  			wait.ForLog("database system is ready to accept connections").
    28  				WithOccurrence(2).
    29  				WithStartupTimeout(5*time.Second)),
    30  	)
    31  
    32  	tearDown := func() {
    33  		_ = pgContainer.Terminate(TestContext)
    34  	}
    35  
    36  	return pgContainer, tearDown, err
    37  }
    38  
    39  // SetupTimescaleContainer creates a PostgreSQL container with the TimescaleDB extension available.
    40  func SetupTimescaleContainer() (*postgres.PostgresContainer, func(), error) {
    41  	pgContainer, err := postgres.Run(TestContext,
    42  		TimescaleImage,
    43  		postgres.WithDatabase(MockDatabase),
    44  		testcontainers.WithWaitStrategy(
    45  			wait.ForLog("database system is ready to accept connections").
    46  				WithOccurrence(2).
    47  				WithStartupTimeout(30*time.Second)),
    48  	)
    49  
    50  	tearDown := func() {
    51  		_ = pgContainer.Terminate(TestContext)
    52  	}
    53  
    54  	return pgContainer, tearDown, err
    55  }
    56  
    57  // Creates a PostgreSQL container with CSV logging enabled.
    58  // This is useful for testing log parsing functionality with server_log_event_counts metric.
    59  func SetupPostgresContainerWithConfig(configPath string) (*postgres.PostgresContainer, func(), error) {
    60  	pgContainer, err := postgres.Run(TestContext,
    61  		PostgresImage,
    62  		postgres.WithDatabase(MockDatabase),
    63  		postgres.WithConfigFile(configPath),
    64  		testcontainers.WithWaitStrategy(
    65  			wait.ForListeningPort("5432/tcp").WithStartupTimeout(5*time.Second)),
    66  	)
    67  
    68  	tearDown := func() {
    69  		_ = pgContainer.Terminate(TestContext)
    70  	}
    71  
    72  	return pgContainer, tearDown, err
    73  }
    74  
    75  func SetupPostgresContainerWithInitScripts(scripts ...string) (*postgres.PostgresContainer, func(), error) {
    76  	pgContainer, err := postgres.Run(TestContext,
    77  		PostgresImage,
    78  		postgres.WithDatabase(MockDatabase),
    79  		postgres.WithInitScripts(scripts...),
    80  		testcontainers.WithWaitStrategy(
    81  			wait.ForLog("database system is ready to accept connections").
    82  				WithOccurrence(2).
    83  				WithStartupTimeout(5*time.Second)),
    84  	)
    85  
    86  	tearDown := func() {
    87  		_ = pgContainer.Terminate(TestContext)
    88  	}
    89  
    90  	return pgContainer, tearDown, err
    91  }
    92  
    93  func SetupEtcdContainer() (*etcd.EtcdContainer, func(), error) {
    94  	etcdContainer, err := etcd.Run(TestContext, EtcdImage,
    95  		testcontainers.
    96  			WithWaitStrategy(wait.ForLog("ready to serve client requests").
    97  				WithStartupTimeout(15*time.Second)))
    98  
    99  	tearDown := func() {
   100  		_ = etcdContainer.Terminate(TestContext)
   101  	}
   102  
   103  	return etcdContainer, tearDown, err
   104  }
   105  
   106  //-----------Setup gRPC test servers-----------------
   107  
   108  func LoadServerTLSCredentials() (credentials.TransportCredentials, error) {
   109  	cert, err := tls.X509KeyPair(Cert, PrivateKey)
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	tlsConfig := &tls.Config{
   115  		Certificates: []tls.Certificate{cert},
   116  	}
   117  	return credentials.NewTLS(tlsConfig), nil
   118  }
   119  
   120  func AuthInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
   121  	md, _ := metadata.FromIncomingContext(ctx)
   122  
   123  	clientUsername := md.Get("username")[0]
   124  	clientPassword := md.Get("password")[0]
   125  
   126  	if clientUsername != "" && clientUsername != "pgwatch" && clientPassword != "pgwatch" {
   127  		return nil, status.Error(codes.Unauthenticated, "unauthenticated")
   128  	}
   129  
   130  	return handler(ctx, req)
   131  }
   132  
   133  func SetupRPCServers() (func(), error) {
   134  	err := os.WriteFile(CAFile, []byte(CA), 0644)
   135  	teardown := func() { _ = os.Remove(CAFile) }
   136  	if err != nil {
   137  		return teardown, err
   138  	}
   139  
   140  	addresses := [2]string{PlainServerAddress, TLSServerAddress}
   141  	for _, address := range addresses {
   142  		lis, err := net.Listen("tcp", address)
   143  		if err != nil {
   144  			return teardown, err
   145  		}
   146  
   147  		var creds credentials.TransportCredentials
   148  		if address == TLSServerAddress {
   149  			creds, err = LoadServerTLSCredentials()
   150  			if err != nil {
   151  				return nil, err
   152  			}
   153  		}
   154  
   155  		server := grpc.NewServer(
   156  			grpc.UnaryInterceptor(AuthInterceptor),
   157  			grpc.Creds(creds),
   158  		)
   159  
   160  		recv := new(Receiver)
   161  		pb.RegisterReceiverServer(server, recv)
   162  
   163  		go func() {
   164  			if err := server.Serve(lis); err != nil {
   165  				panic(err)
   166  			}
   167  		}()
   168  	}
   169  	// wait a little for servers to start
   170  	time.Sleep(time.Second)
   171  	return teardown, nil
   172  }
   173