github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/pusher/v2/tenant_pusher_test.go (about)

     1  package v2
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/golang/snappy"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/stretchr/testify/require"
    16  	"golang.org/x/sync/errgroup"
    17  
    18  	"github.com/grafana/synthetic-monitoring-agent/internal/pusher"
    19  	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
    20  )
    21  
    22  func TestTenantPusher(t *testing.T) {
    23  	// This is an extremely basic test that verifies that a tenant pusher
    24  	// can be constructed.
    25  	tenantProvider := testTenantProvider{
    26  		1: &sm.Tenant{
    27  			Id:            1,
    28  			OrgId:         1,
    29  			MetricsRemote: &sm.RemoteInfo{},
    30  			EventsRemote:  &sm.RemoteInfo{},
    31  			Status:        sm.TenantStatus_ACTIVE,
    32  		},
    33  	}
    34  
    35  	registry := prometheus.NewPedanticRegistry()
    36  	metrics := pusher.NewMetrics(registry)
    37  
    38  	p := newTenantPusher(1, tenantProvider, pusherOptions{
    39  		metrics: metrics,
    40  	})
    41  
    42  	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(50*time.Millisecond))
    43  	defer cancel()
    44  
    45  	var wg sync.WaitGroup
    46  	wg.Add(1)
    47  
    48  	go func() {
    49  		p.run(ctx)
    50  		wg.Done()
    51  	}()
    52  
    53  	wg.Wait()
    54  }
    55  
    56  func TestTenantPusherRunPushers(t *testing.T) {
    57  	// This test will make sure that runPushers handles error conditions correctly.
    58  	tenantProvider := testTenantProvider{
    59  		1: &sm.Tenant{
    60  			Id:            1,
    61  			OrgId:         1,
    62  			MetricsRemote: &sm.RemoteInfo{},
    63  			EventsRemote:  &sm.RemoteInfo{},
    64  			Status:        sm.TenantStatus_ACTIVE,
    65  		},
    66  	}
    67  
    68  	registry := prometheus.NewPedanticRegistry()
    69  	metrics := pusher.NewMetrics(registry).WithTenant(2, 1)
    70  
    71  	// use tenant ID 2 to force an error
    72  	p := newTenantPusher(2, tenantProvider, pusherOptions{
    73  		metrics: metrics,
    74  	})
    75  
    76  	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(50*time.Millisecond))
    77  	defer cancel()
    78  
    79  	g, gCtx := errgroup.WithContext(ctx)
    80  
    81  	g.Go(func() error {
    82  		// Make sure the deadline is handled. If the runPushers method
    83  		// below returns before the deadline, the group context will be
    84  		// cancelled and the line below will return. If the deadline is
    85  		// met, the line below will return instead (if runPushers
    86  		// doesn't handle it correctly).
    87  		<-gCtx.Done()
    88  		return gCtx.Err()
    89  	})
    90  
    91  	g.Go(func() error { return p.runPushers(gCtx) })
    92  
    93  	err := g.Wait()
    94  
    95  	require.Error(t, err)
    96  	require.ErrorIs(t, err, errTestNoTenant)
    97  }
    98  
    99  func makeRecords(blocks [][]byte) []queueEntry {
   100  	out := make([]queueEntry, len(blocks))
   101  	for idx, b := range blocks {
   102  		data := make([]byte, len(b))
   103  		copy(data, b)
   104  		out[idx].data = &data
   105  	}
   106  
   107  	return out
   108  }
   109  
   110  type testTenantProvider map[int64]*sm.Tenant
   111  
   112  var errTestNoTenant = errors.New("tenant not found")
   113  
   114  func (t testTenantProvider) GetTenant(ctx context.Context, info *sm.TenantInfo) (*sm.Tenant, error) {
   115  	tenant, found := t[info.Id]
   116  	if !found {
   117  		return nil, errTestNoTenant
   118  	}
   119  
   120  	return tenant, nil
   121  }
   122  
   123  type testServer struct {
   124  	mu           sync.Mutex
   125  	receivedBody []byte
   126  	responses    []http.HandlerFunc
   127  	server       *httptest.Server
   128  }
   129  
   130  func (s *testServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   131  	s.mu.Lock()
   132  	defer s.mu.Unlock()
   133  
   134  	if len(s.responses) == 0 {
   135  		panic(len(s.responses))
   136  	}
   137  
   138  	body, err := io.ReadAll(r.Body)
   139  	if err != nil {
   140  		panic(err)
   141  	}
   142  
   143  	content, err := snappy.Decode(nil, body)
   144  	if err != nil {
   145  		panic(err)
   146  	}
   147  
   148  	s.receivedBody = content
   149  
   150  	defer r.Body.Close()
   151  
   152  	act := s.responses[0]
   153  	s.responses = s.responses[1:]
   154  
   155  	act(w, r)
   156  }
   157  
   158  func (s *testServer) start() {
   159  	s.server = httptest.NewServer(s)
   160  }
   161  
   162  func (s *testServer) stop() {
   163  	s.server.Close()
   164  }
   165  
   166  func (s *testServer) done() bool {
   167  	s.mu.Lock()
   168  	defer s.mu.Unlock()
   169  
   170  	return len(s.responses) == 0
   171  }