mocks.go

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

     1  package testutil
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"github.com/cybertec-postgresql/pgwatch/v6/api/pb"
     8  	"github.com/cybertec-postgresql/pgwatch/v6/internal/db"
     9  	"github.com/cybertec-postgresql/pgwatch/v6/internal/metrics"
    10  	"github.com/cybertec-postgresql/pgwatch/v6/internal/sources"
    11  	"github.com/jackc/pgx/v5"
    12  	"github.com/jackc/pgx/v5/pgconn"
    13  	"github.com/jackc/pgx/v5/pgxpool"
    14  	"google.golang.org/protobuf/types/known/structpb"
    15  )
    16  
    17  // Receiver implements the ReceiverServer interface for testing purposes
    18  type Receiver struct {
    19  	pb.UnimplementedReceiverServer
    20  }
    21  
    22  func (receiver *Receiver) UpdateMeasurements(_ context.Context, msg *pb.MeasurementEnvelope) (*pb.Reply, error) {
    23  	if len(msg.GetData()) == 0 {
    24  		return nil, errors.New("empty message")
    25  	}
    26  	if msg.GetDBName() != "Db" {
    27  		return nil, errors.New("invalid message")
    28  	}
    29  	return &pb.Reply{}, nil
    30  }
    31  
    32  func (receiver *Receiver) SyncMetric(_ context.Context, syncReq *pb.SyncReq) (*pb.Reply, error) {
    33  	if syncReq == nil {
    34  		return nil, errors.New("nil sync request")
    35  	}
    36  	if syncReq.GetOperation() == pb.SyncOp_InvalidOp {
    37  		return nil, errors.New("invalid sync request")
    38  	}
    39  	return &pb.Reply{}, nil
    40  }
    41  
    42  func (receiver *Receiver) DefineMetrics(_ context.Context, metricsStruct *structpb.Struct) (*pb.Reply, error) {
    43  	if metricsStruct == nil {
    44  		return nil, errors.New("nil metrics struct")
    45  	}
    46  	if metricsStruct.GetFields() == nil {
    47  		return nil, errors.New("empty metrics struct")
    48  	}
    49  	return &pb.Reply{Logmsg: "metrics defined successfully"}, nil
    50  }
    51  
    52  // MockMetricsReaderWriter implements MetricsReaderWriter interface
    53  type MockMetricsReaderWriter struct {
    54  	GetMetricsFunc   func() (*metrics.Metrics, error)
    55  	UpdateMetricFunc func(name string, m metrics.Metric) error
    56  	CreateMetricFunc func(name string, m metrics.Metric) error
    57  	DeleteMetricFunc func(name string) error
    58  	DeletePresetFunc func(name string) error
    59  	UpdatePresetFunc func(name string, preset metrics.Preset) error
    60  	CreatePresetFunc func(name string, preset metrics.Preset) error
    61  	WriteMetricsFunc func(metricDefs *metrics.Metrics) error
    62  }
    63  
    64  func (m *MockMetricsReaderWriter) GetMetrics() (*metrics.Metrics, error) {
    65  	return m.GetMetricsFunc()
    66  }
    67  func (m *MockMetricsReaderWriter) UpdateMetric(name string, metric metrics.Metric) error {
    68  	return m.UpdateMetricFunc(name, metric)
    69  }
    70  func (m *MockMetricsReaderWriter) CreateMetric(name string, metric metrics.Metric) error {
    71  	return m.CreateMetricFunc(name, metric)
    72  }
    73  func (m *MockMetricsReaderWriter) DeleteMetric(name string) error {
    74  	return m.DeleteMetricFunc(name)
    75  }
    76  func (m *MockMetricsReaderWriter) DeletePreset(name string) error {
    77  	return m.DeletePresetFunc(name)
    78  }
    79  func (m *MockMetricsReaderWriter) UpdatePreset(name string, preset metrics.Preset) error {
    80  	return m.UpdatePresetFunc(name, preset)
    81  }
    82  func (m *MockMetricsReaderWriter) CreatePreset(name string, preset metrics.Preset) error {
    83  	return m.CreatePresetFunc(name, preset)
    84  }
    85  func (m *MockMetricsReaderWriter) WriteMetrics(metricDefs *metrics.Metrics) error {
    86  	return m.WriteMetricsFunc(metricDefs)
    87  }
    88  
    89  // MockSourcesReaderWriter implements SourcesReaderWriter interface
    90  type MockSourcesReaderWriter struct {
    91  	GetSourcesFunc   func() (sources.Sources, error)
    92  	UpdateSourceFunc func(md sources.Source) error
    93  	CreateSourceFunc func(md sources.Source) error
    94  	DeleteSourceFunc func(name string) error
    95  	WriteSourcesFunc func(sources.Sources) error
    96  }
    97  
    98  func (m *MockSourcesReaderWriter) GetSources() (sources.Sources, error) {
    99  	return m.GetSourcesFunc()
   100  }
   101  func (m *MockSourcesReaderWriter) UpdateSource(md sources.Source) error {
   102  	return m.UpdateSourceFunc(md)
   103  }
   104  func (m *MockSourcesReaderWriter) CreateSource(md sources.Source) error {
   105  	return m.CreateSourceFunc(md)
   106  }
   107  func (m *MockSourcesReaderWriter) DeleteSource(name string) error {
   108  	return m.DeleteSourceFunc(name)
   109  }
   110  func (m *MockSourcesReaderWriter) WriteSources(srcs sources.Sources) error {
   111  	return m.WriteSourcesFunc(srcs)
   112  }
   113  
   114  // BlockingPool is a wedged PgxPoolIface used by fault-injection tests.
   115  //
   116  // It embeds db.PgxPoolIface as a nil interface so the type satisfies
   117  // db.PgxPoolIface; only the methods overridden below are usable. Any
   118  // direct call to a non-overridden method panics (nil interface call).
   119  //
   120  // Ping, Query, SendBatch, and Acquire all block until the supplied
   121  // context is cancelled, then return ctx.Err(). This simulates a pool
   122  // whose connections are stuck on the wire — the failure mode that
   123  // client-side deadlines must convert into bounded failures.
   124  type BlockingPool struct {
   125  	db.PgxPoolIface
   126  }
   127  
   128  // Ping blocks until ctx.Done() and returns ctx.Err().
   129  func (BlockingPool) Ping(ctx context.Context) error {
   130  	<-ctx.Done()
   131  	return ctx.Err()
   132  }
   133  
   134  // Query blocks until ctx.Done() and returns ctx.Err().
   135  func (BlockingPool) Query(ctx context.Context, _ string, _ ...any) (pgx.Rows, error) {
   136  	<-ctx.Done()
   137  	return nil, ctx.Err()
   138  }
   139  
   140  // QueryRow blocks until ctx.Done() and returns a row whose Scan returns ctx.Err().
   141  func (BlockingPool) QueryRow(ctx context.Context, _ string, _ ...any) pgx.Row {
   142  	return blockingRow{ctx: ctx}
   143  }
   144  
   145  // Exec blocks until ctx.Done() and returns ctx.Err().
   146  func (BlockingPool) Exec(ctx context.Context, _ string, _ ...any) (pgconn.CommandTag, error) {
   147  	<-ctx.Done()
   148  	return pgconn.CommandTag{}, ctx.Err()
   149  }
   150  
   151  // first Query call returns ctx.Err().
   152  // SendBatch blocks until ctx.Done() and returns a BatchResults whose
   153  // first Query call returns ctx.Err().
   154  func (BlockingPool) SendBatch(ctx context.Context, _ *pgx.Batch) pgx.BatchResults {
   155  	return &BlockingBatchResults{ctx: ctx}
   156  }
   157  
   158  // Acquire blocks until ctx.Done() and returns ctx.Err().
   159  func (BlockingPool) Acquire(ctx context.Context) (*pgxpool.Conn, error) {
   160  	<-ctx.Done()
   161  	return nil, ctx.Err()
   162  }
   163  
   164  // Close is a no-op so the pool can be embedded without panicking.
   165  func (BlockingPool) Close() {}
   166  
   167  // blockingRow makes pgx.Row.Scan honor ctx cancellation.
   168  type blockingRow struct {
   169  	ctx context.Context
   170  }
   171  
   172  // Scan blocks until ctx.Done() and returns ctx.Err().
   173  func (b blockingRow) Scan(_ ...any) error {
   174  	<-b.ctx.Done()
   175  	return b.ctx.Err()
   176  }
   177  
   178  // BlockingBatchResults makes the first Query call honor ctx cancellation.
   179  type BlockingBatchResults struct {
   180  	ctx    context.Context
   181  	closed bool
   182  }
   183  
   184  // Query blocks until ctx.Done() and returns ctx.Err().
   185  func (b *BlockingBatchResults) Query() (pgx.Rows, error) {
   186  	<-b.ctx.Done()
   187  	return nil, b.ctx.Err()
   188  }
   189  
   190  // Exec blocks until ctx.Done() and returns ctx.Err().
   191  func (b *BlockingBatchResults) Exec() (pgconn.CommandTag, error) {
   192  	<-b.ctx.Done()
   193  	return pgconn.CommandTag{}, b.ctx.Err()
   194  }
   195  func (b *BlockingBatchResults) Close() error {
   196  	b.closed = true
   197  	return nil
   198  }
   199  
   200  // Err returns ctx.Err() if the batch was closed via Close.
   201  
   202  // QueryRow blocks until ctx.Done() and returns a Row whose Scan returns ctx.Err().
   203  func (b *BlockingBatchResults) QueryRow() pgx.Row {
   204  	return blockingRow{ctx: b.ctx}
   205  }
   206  func (b *BlockingBatchResults) Err() error {
   207  	if b.closed {
   208  		return b.ctx.Err()
   209  	}
   210  	return nil
   211  }
   212