github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/checks/checks_test.go (about)

     1  package checks
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync/atomic"
     7  	"syscall"
     8  	"testing"
     9  	"testing/synctest"
    10  	"time"
    11  
    12  	"github.com/grafana/synthetic-monitoring-agent/internal/secrets"
    13  
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/prometheus/client_golang/prometheus/testutil"
    16  	"github.com/rs/zerolog"
    17  	"github.com/stretchr/testify/require"
    18  	"google.golang.org/grpc"
    19  
    20  	"github.com/grafana/synthetic-monitoring-agent/internal/feature"
    21  	"github.com/grafana/synthetic-monitoring-agent/internal/k6runner"
    22  	"github.com/grafana/synthetic-monitoring-agent/internal/model"
    23  	"github.com/grafana/synthetic-monitoring-agent/internal/prober"
    24  	"github.com/grafana/synthetic-monitoring-agent/internal/prober/logger"
    25  	"github.com/grafana/synthetic-monitoring-agent/internal/pusher"
    26  	"github.com/grafana/synthetic-monitoring-agent/internal/scraper"
    27  	"github.com/grafana/synthetic-monitoring-agent/internal/telemetry"
    28  	"github.com/grafana/synthetic-monitoring-agent/internal/testhelper"
    29  	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
    30  )
    31  
    32  func TestNewUpdater(t *testing.T) {
    33  	testFeatureCollection := feature.NewCollection()
    34  	require.NotNil(t, testFeatureCollection)
    35  	require.NoError(t, testFeatureCollection.Set("foo"))
    36  	require.NoError(t, testFeatureCollection.Set("bar"))
    37  
    38  	testcases := map[string]struct {
    39  		opts UpdaterOptions
    40  	}{
    41  		"trivial": {
    42  			opts: UpdaterOptions{
    43  				Conn:           new(grpc.ClientConn),
    44  				PromRegisterer: prometheus.NewPedanticRegistry(),
    45  				Publisher:      channelPublisher(make(chan pusher.Payload)),
    46  				TenantCh:       make(chan<- sm.Tenant),
    47  				Logger:         testhelper.Logger(t),
    48  				Features:       testFeatureCollection,
    49  			},
    50  		},
    51  	}
    52  
    53  	for name, tc := range testcases {
    54  		t.Run(name, func(t *testing.T) {
    55  			u, err := NewUpdater(tc.opts)
    56  			require.NoError(t, err)
    57  			require.NotNil(t, u)
    58  			require.Equal(t, tc.opts.Publisher, u.publisher)
    59  			require.Equal(t, tc.opts.TenantCh, u.tenantCh)
    60  			require.Equal(t, tc.opts.Features, u.features)
    61  			require.Equal(t, tc.opts.Logger, u.logger)
    62  			require.Equal(t, tc.opts.Conn, u.api.conn)
    63  			require.NotNil(t, u.scrapers)
    64  			require.NotNil(t, u.metrics.changesCounter)
    65  			require.NotNil(t, u.metrics.changeErrorsCounter)
    66  			require.NotNil(t, u.metrics.runningScrapers)
    67  			require.NotNil(t, u.metrics.scrapesCounter)
    68  			require.NotNil(t, u.metrics.scrapeErrorCounter)
    69  			require.NotNil(t, u.metrics.probeInfo)
    70  			require.False(t, u.supportsProtocolSecrets, "default value should be false")
    71  		})
    72  	}
    73  }
    74  
    75  func TestNewUpdaterSupportsProtocolSecrets(t *testing.T) {
    76  	testFeatureCollection := feature.NewCollection()
    77  	require.NotNil(t, testFeatureCollection)
    78  
    79  	opts := UpdaterOptions{
    80  		Conn:                    new(grpc.ClientConn),
    81  		PromRegisterer:          prometheus.NewPedanticRegistry(),
    82  		Publisher:               channelPublisher(make(chan pusher.Payload)),
    83  		TenantCh:                make(chan<- sm.Tenant),
    84  		Logger:                  testhelper.Logger(t),
    85  		Features:                testFeatureCollection,
    86  		SupportsProtocolSecrets: true,
    87  	}
    88  
    89  	u, err := NewUpdater(opts)
    90  	require.NoError(t, err)
    91  	require.NotNil(t, u)
    92  	require.True(t, u.supportsProtocolSecrets, "should be set to true")
    93  }
    94  
    95  func TestInstallSignalHandler(t *testing.T) {
    96  	testcases := map[string]func(t *testing.T){
    97  		"signal": func(t *testing.T) {
    98  			// verify that the signal context is done after
    99  			// receiving the signal, and that the signal is
   100  			// correctly reported as having fired.
   101  			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
   102  			defer cancel()
   103  
   104  			sigCtx, signalFired := installSignalHandler(ctx)
   105  			require.NotNil(t, sigCtx)
   106  			require.NotNil(t, signalFired)
   107  			require.NoError(t, syscall.Kill(syscall.Getpid(), syscall.SIGUSR1))
   108  
   109  			select {
   110  			case <-ctx.Done():
   111  				t.Fatal("context timeout expired")
   112  			case <-sigCtx.Done():
   113  				require.Equal(t, int32(1), atomic.LoadInt32(signalFired))
   114  			}
   115  		},
   116  
   117  		"no signal": func(t *testing.T) {
   118  			// verify that the signal context is done after
   119  			// the parrent context is done, and that the
   120  			// signal is correctly reported as not having
   121  			// fired.
   122  			ctx, cancel := context.WithCancel(context.Background())
   123  			sigCtx, signalFired := installSignalHandler(ctx)
   124  			require.NotNil(t, sigCtx)
   125  			require.NotNil(t, signalFired)
   126  
   127  			cancel()
   128  
   129  			timeout := 100 * time.Millisecond
   130  
   131  			timer := time.NewTimer(timeout)
   132  			defer timer.Stop()
   133  
   134  			select {
   135  			case <-timer.C:
   136  				t.Fatalf("signal context not cancelled after %s", timeout)
   137  			case <-sigCtx.Done():
   138  				require.Equal(t, int32(0), atomic.LoadInt32(signalFired))
   139  			}
   140  		},
   141  	}
   142  
   143  	for name, f := range testcases {
   144  		t.Run(name, f)
   145  	}
   146  }
   147  
   148  func TestSleepCtx(t *testing.T) {
   149  	var (
   150  		veryShort = 1 * time.Microsecond
   151  		long      = 10 * time.Second
   152  	)
   153  
   154  	// make sure errors are reported correctly
   155  
   156  	ctx := context.Background()
   157  	err := sleepCtx(ctx, veryShort)
   158  	require.NoError(t, err)
   159  
   160  	ctx, cancel := context.WithCancel(context.Background())
   161  	cancel()
   162  
   163  	err = sleepCtx(ctx, long)
   164  	require.Error(t, err)
   165  	require.ErrorIs(t, err, context.Canceled)
   166  
   167  	ctx, cancel = context.WithTimeout(context.Background(), veryShort)
   168  	err = sleepCtx(ctx, long)
   169  	require.Error(t, err)
   170  	require.ErrorIs(t, err, context.DeadlineExceeded)
   171  	cancel()
   172  
   173  	ctx, cancel = context.WithTimeout(context.Background(), long)
   174  	cancel()
   175  
   176  	err = sleepCtx(ctx, long)
   177  	require.Error(t, err)
   178  	require.ErrorIs(t, err, context.Canceled)
   179  }
   180  
   181  // TestHandleCheckOp is testing internal functions that run as part of
   182  // updater.Run. Since these functions operate on scraper instances, a
   183  // test scraper is used, which in turn creates a test probe. The goal of
   184  // this is to decouple the testing of these functions from the testing
   185  // of the prober themselves.
   186  func TestHandleCheckOp(t *testing.T) {
   187  	synctest.Test(t, testHandleCheckOpImpl)
   188  }
   189  
   190  // testHandleCheckOpImpl is the actual implementation of the
   191  // TestHandleCheckOp test above to avoid excessive indentation.
   192  func testHandleCheckOpImpl(t *testing.T) {
   193  	publishCh := make(chan pusher.Payload, 100)
   194  
   195  	u, err := NewUpdater(
   196  		UpdaterOptions{
   197  			Conn:           new(grpc.ClientConn),
   198  			PromRegisterer: prometheus.NewPedanticRegistry(),
   199  			Publisher:      channelPublisher(publishCh),
   200  			TenantCh:       make(chan<- sm.Tenant),
   201  			Logger:         testhelper.Logger(t),
   202  			ScraperFactory: testScraperFactory,
   203  		},
   204  	)
   205  
   206  	require.NotNil(t, u)
   207  	require.NoError(t, err)
   208  
   209  	u.probe = &sm.Probe{
   210  		Id:   100,
   211  		Name: "test-probe",
   212  	}
   213  
   214  	// Since this is meant to run inside a test bubble, t.Deadline cannot be used.
   215  	// The timeout here doesn't really matter because this is running with faketime.
   216  	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
   217  	defer cancel()
   218  
   219  	var check model.Check
   220  
   221  	err = check.FromSM(sm.Check{
   222  		Id:        5000,
   223  		TenantId:  1,
   224  		Frequency: 1000,
   225  		Timeout:   1000,
   226  		Target:    "127.0.0.1",
   227  		Job:       "", // not setting value to make check invalid
   228  		Probes:    []int64{1},
   229  		Settings: sm.CheckSettings{
   230  			Ping: &sm.PingSettings{},
   231  		},
   232  		Created:  0,
   233  		Modified: 0,
   234  	})
   235  	require.NoError(t, err)
   236  
   237  	scraperExists := func() bool {
   238  		u.scrapersMutex.Lock()
   239  		defer u.scrapersMutex.Unlock()
   240  
   241  		_, found := u.scrapers[check.GlobalID()]
   242  
   243  		return found
   244  	}
   245  
   246  	// this should fail, check is invalid
   247  	err = u.handleCheckAdd(ctx, check)
   248  	require.Error(t, err)
   249  	// This doesn't work because the counter hasn't been set
   250  	// (because of the error):
   251  	// require.Equal(t, 0.0, testutil.ToFloat64(u.metrics.runningScrapers))
   252  	require.False(t, scraperExists())
   253  
   254  	// fix check
   255  	check.Job = "test-job"
   256  	check.Modified++
   257  
   258  	err = u.handleCheckAdd(ctx, check)
   259  	require.NoError(t, err)
   260  	require.Equal(t, 1.0, testutil.ToFloat64(u.metrics.runningScrapers))
   261  	require.True(t, scraperExists())
   262  
   263  	check.Modified++
   264  
   265  	// try to add again, this should fail, even if modified changed
   266  	err = u.handleCheckAdd(ctx, check)
   267  	require.Error(t, err)
   268  	require.Equal(t, 1.0, testutil.ToFloat64(u.metrics.runningScrapers))
   269  	require.True(t, scraperExists())
   270  
   271  	check.Modified++
   272  
   273  	// update the existing check
   274  	err = u.handleCheckUpdate(ctx, check)
   275  	require.NoError(t, err)
   276  	require.Equal(t, 1.0, testutil.ToFloat64(u.metrics.runningScrapers))
   277  	require.True(t, scraperExists())
   278  
   279  	err = u.handleCheckDelete(ctx, check)
   280  	require.NoError(t, err)
   281  	require.Equal(t, 0.0, testutil.ToFloat64(u.metrics.runningScrapers))
   282  	require.False(t, scraperExists())
   283  
   284  	// try to delete again
   285  	err = u.handleCheckDelete(ctx, check)
   286  	require.Error(t, err)
   287  	require.Equal(t, 0.0, testutil.ToFloat64(u.metrics.runningScrapers))
   288  	require.False(t, scraperExists())
   289  
   290  	// updating a non-existing check becomes an add
   291  	err = u.handleCheckUpdate(ctx, check)
   292  	require.NoError(t, err)
   293  	require.Equal(t, 1.0, testutil.ToFloat64(u.metrics.runningScrapers))
   294  	require.True(t, scraperExists())
   295  
   296  	// clean up
   297  	err = u.handleCheckDelete(ctx, check)
   298  	require.NoError(t, err)
   299  	require.Equal(t, 0.0, testutil.ToFloat64(u.metrics.runningScrapers))
   300  	require.False(t, scraperExists())
   301  
   302  	// Wait for all scraper goroutines to fully exit before test completes.
   303  	synctest.Wait()
   304  }
   305  
   306  func TestCheckHandlerProbeValidation(t *testing.T) {
   307  	t.Parallel()
   308  
   309  	testcases := map[string]struct {
   310  		opts          UpdaterOptions
   311  		probe         sm.Probe
   312  		expectedError error
   313  	}{
   314  		"has K6 when required for scripted checks": {
   315  			expectedError: nil,
   316  			opts: UpdaterOptions{
   317  				Conn:           new(grpc.ClientConn),
   318  				PromRegisterer: prometheus.NewPedanticRegistry(),
   319  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   320  				TenantCh:       make(chan<- sm.Tenant),
   321  				Logger:         testhelper.Logger(t),
   322  				K6Runner:       noopRunner{},
   323  			},
   324  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   325  				DisableScriptedChecks: false,
   326  				DisableBrowserChecks:  true,
   327  			}},
   328  		},
   329  		"missing K6 when required for scripted checks": {
   330  			expectedError: errCapabilityK6Missing,
   331  			opts: UpdaterOptions{
   332  				Conn:           new(grpc.ClientConn),
   333  				PromRegisterer: prometheus.NewPedanticRegistry(),
   334  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   335  				TenantCh:       make(chan<- sm.Tenant),
   336  				Logger:         testhelper.Logger(t),
   337  			},
   338  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   339  				DisableScriptedChecks: false,
   340  				DisableBrowserChecks:  true,
   341  			}},
   342  		},
   343  		"has K6 when required for browser checks": {
   344  			expectedError: nil,
   345  			opts: UpdaterOptions{
   346  				Conn:           new(grpc.ClientConn),
   347  				PromRegisterer: prometheus.NewPedanticRegistry(),
   348  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   349  				TenantCh:       make(chan<- sm.Tenant),
   350  				Logger:         testhelper.Logger(t),
   351  				K6Runner:       noopRunner{},
   352  			},
   353  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   354  				DisableScriptedChecks: true,
   355  				DisableBrowserChecks:  false,
   356  			}},
   357  		},
   358  		"missing K6 when required for browser checks": {
   359  			expectedError: errCapabilityK6Missing,
   360  			opts: UpdaterOptions{
   361  				Conn:           new(grpc.ClientConn),
   362  				PromRegisterer: prometheus.NewPedanticRegistry(),
   363  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   364  				TenantCh:       make(chan<- sm.Tenant),
   365  				Logger:         testhelper.Logger(t),
   366  			},
   367  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   368  				DisableScriptedChecks: true,
   369  				DisableBrowserChecks:  false,
   370  			}},
   371  		},
   372  		"has K6 when required for scripted and browser checks": {
   373  			expectedError: nil,
   374  			opts: UpdaterOptions{
   375  				Conn:           new(grpc.ClientConn),
   376  				PromRegisterer: prometheus.NewPedanticRegistry(),
   377  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   378  				TenantCh:       make(chan<- sm.Tenant),
   379  				Logger:         testhelper.Logger(t),
   380  				K6Runner:       noopRunner{},
   381  			},
   382  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   383  				DisableScriptedChecks: false,
   384  				DisableBrowserChecks:  false,
   385  			}},
   386  		},
   387  		"missing K6 when required for scripted and browser checks": {
   388  			expectedError: errCapabilityK6Missing,
   389  			opts: UpdaterOptions{
   390  				Conn:           new(grpc.ClientConn),
   391  				PromRegisterer: prometheus.NewPedanticRegistry(),
   392  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   393  				TenantCh:       make(chan<- sm.Tenant),
   394  				Logger:         testhelper.Logger(t),
   395  			},
   396  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   397  				DisableScriptedChecks: false,
   398  				DisableBrowserChecks:  false,
   399  			}},
   400  		},
   401  		"has K6 but not required": {
   402  			expectedError: nil,
   403  			opts: UpdaterOptions{
   404  				Conn:           new(grpc.ClientConn),
   405  				PromRegisterer: prometheus.NewPedanticRegistry(),
   406  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   407  				TenantCh:       make(chan<- sm.Tenant),
   408  				Logger:         testhelper.Logger(t),
   409  				K6Runner:       noopRunner{},
   410  			},
   411  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   412  				DisableScriptedChecks: true,
   413  				DisableBrowserChecks:  true,
   414  			}},
   415  		},
   416  		"missing K6 but not required": {
   417  			expectedError: nil,
   418  			opts: UpdaterOptions{
   419  				Conn:           new(grpc.ClientConn),
   420  				PromRegisterer: prometheus.NewPedanticRegistry(),
   421  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   422  				TenantCh:       make(chan<- sm.Tenant),
   423  				Logger:         testhelper.Logger(t),
   424  			},
   425  			probe: sm.Probe{Id: 100, Name: "test-probe", Capabilities: &sm.Probe_Capabilities{
   426  				DisableScriptedChecks: true,
   427  				DisableBrowserChecks:  true,
   428  			}},
   429  		},
   430  		"missing K6 when required by default": {
   431  			expectedError: errCapabilityK6Missing,
   432  			opts: UpdaterOptions{
   433  				Conn:           new(grpc.ClientConn),
   434  				PromRegisterer: prometheus.NewPedanticRegistry(),
   435  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   436  				TenantCh:       make(chan<- sm.Tenant),
   437  				Logger:         testhelper.Logger(t),
   438  			},
   439  			probe: sm.Probe{Id: 100, Name: "test-probe"},
   440  		},
   441  		"has K6 when required by default": {
   442  			expectedError: nil,
   443  			opts: UpdaterOptions{
   444  				Conn:           new(grpc.ClientConn),
   445  				PromRegisterer: prometheus.NewPedanticRegistry(),
   446  				Publisher:      channelPublisher(make(chan pusher.Payload)),
   447  				TenantCh:       make(chan<- sm.Tenant),
   448  				Logger:         testhelper.Logger(t),
   449  				K6Runner:       noopRunner{},
   450  			},
   451  			probe: sm.Probe{Id: 100, Name: "test-probe"},
   452  		},
   453  	}
   454  
   455  	for testName, tc := range testcases {
   456  		t.Run(testName, func(t *testing.T) {
   457  			t.Parallel()
   458  
   459  			u, err := NewUpdater(tc.opts)
   460  			require.NoError(t, err)
   461  
   462  			err = u.validateProbeCapabilities(tc.probe.Capabilities)
   463  
   464  			if tc.expectedError != nil {
   465  				require.Error(t, err, tc.expectedError)
   466  			} else {
   467  				require.NoError(t, err)
   468  			}
   469  		})
   470  	}
   471  }
   472  
   473  type testProber struct{}
   474  
   475  func (testProber) Name() string {
   476  	return "test-prober"
   477  }
   478  
   479  func (testProber) Probe(ctx context.Context, target string, registry *prometheus.Registry, logger logger.Logger, _ string) (bool, float64) {
   480  	return false, 0
   481  }
   482  
   483  type testProbeFactory struct{}
   484  
   485  func (f testProbeFactory) New(ctx context.Context, logger zerolog.Logger, check model.Check) (prober.Prober, string, error) {
   486  	return testProber{}, check.Target, nil
   487  }
   488  
   489  type testLabelsLimiter struct {
   490  	metricLabelsLimit int
   491  	logLabelsLimit    int
   492  }
   493  
   494  func (l testLabelsLimiter) MetricLabels(ctx context.Context, tenantID model.GlobalID) (int, error) {
   495  	return l.metricLabelsLimit, nil
   496  }
   497  
   498  func (l testLabelsLimiter) LogLabels(ctx context.Context, tenantID model.GlobalID) (int, error) {
   499  	return l.logLabelsLimit, nil
   500  }
   501  
   502  type testLabellingMode struct{}
   503  
   504  func (l testLabellingMode) ForTenant(ctx context.Context, tenantID model.GlobalID) (sm.LabelMode, error) {
   505  	return sm.LabelMode_LABEL_MODE_PREFIXED, nil
   506  }
   507  
   508  func testScraperFactory(ctx context.Context, check model.Check, publisher pusher.Publisher, _ sm.Probe,
   509  	_ feature.Collection,
   510  	logger zerolog.Logger,
   511  	metrics scraper.Metrics,
   512  	k6Runner k6runner.Runner,
   513  	labelsLimiter scraper.LabelsLimiter,
   514  	telemeter *telemetry.Telemeter,
   515  	secretStore secrets.SecretProvider,
   516  	cals scraper.TenantCals,
   517  	labellingMode scraper.TenantLabelMode,
   518  ) (*scraper.Scraper, error) {
   519  	return scraper.NewWithOpts(
   520  		ctx,
   521  		check,
   522  		scraper.ScraperOpts{
   523  			Logger:        logger,
   524  			ProbeFactory:  testProbeFactory{},
   525  			Publisher:     publisher,
   526  			Metrics:       metrics,
   527  			LabelsLimiter: testLabelsLimiter{},
   528  			LabellingMode: testLabellingMode{},
   529  			Telemeter:     telemeter,
   530  		},
   531  	)
   532  }
   533  
   534  var _ scraper.Factory = testScraperFactory
   535  
   536  type channelPublisher chan pusher.Payload
   537  
   538  func (c channelPublisher) Publish(payload pusher.Payload) {
   539  	c <- payload
   540  }
   541  
   542  type noopRunner struct{}
   543  
   544  func (noopRunner) WithLogger(logger *zerolog.Logger) k6runner.Runner {
   545  	var r noopRunner
   546  	return r
   547  }
   548  
   549  func (noopRunner) Run(ctx context.Context, script k6runner.Script, secretStore k6runner.SecretStore, _ string) (*k6runner.RunResponse, error) {
   550  	return &k6runner.RunResponse{}, nil
   551  }
   552  
   553  func (noopRunner) Versions(_ context.Context) <-chan []string {
   554  	return nil // Blocks forever if read.
   555  }
   556  
   557  type testBackoff time.Duration
   558  
   559  func (b *testBackoff) Reset() {
   560  	*b = 0
   561  }
   562  
   563  func (b testBackoff) Duration() time.Duration {
   564  	return time.Duration(b)
   565  }
   566  
   567  // TestHandleError tests the handleError function. It considers the errors that
   568  // might be returned from the loop method.
   569  func TestHandleError(t *testing.T) {
   570  	ctx, cancel := testhelper.Context(context.Background(), t)
   571  	defer cancel()
   572  
   573  	logger := testhelper.Logger(t)
   574  
   575  	t.Run("no error", func(t *testing.T) {
   576  		backoff := testBackoff(1)
   577  		done, err := handleError(ctx, logger, &backoff, false, nil)
   578  		require.True(t, done)
   579  		require.NoError(t, err)
   580  		require.NotZero(t, backoff)
   581  	})
   582  
   583  	t.Run("context cancelled", func(t *testing.T) {
   584  		backoff := testBackoff(1)
   585  		done, err := handleError(ctx, logger, &backoff, false, fmt.Errorf("wrapped: %w", context.Canceled))
   586  		require.True(t, done)
   587  		require.NoError(t, err)
   588  		require.NotZero(t, backoff)
   589  	})
   590  
   591  	t.Run("k6 capability missing", func(t *testing.T) {
   592  		backoff := testBackoff(1)
   593  		done, err := handleError(ctx, logger, &backoff, false, errCapabilityK6Missing)
   594  		require.True(t, done)
   595  		require.ErrorIs(t, err, errCapabilityK6Missing)
   596  		require.NotZero(t, backoff)
   597  	})
   598  
   599  	t.Run("incompatible API", func(t *testing.T) {
   600  		backoff := testBackoff(1)
   601  		done, err := handleError(ctx, logger, &backoff, false, errIncompatibleApi)
   602  		require.True(t, done)
   603  		require.ErrorIs(t, err, errIncompatibleApi)
   604  		require.NotZero(t, backoff)
   605  	})
   606  
   607  	t.Run("not authorized", func(t *testing.T) {
   608  		backoff := testBackoff(1)
   609  		done, err := handleError(ctx, logger, &backoff, false, errNotAuthorized)
   610  		require.True(t, done)
   611  		require.ErrorIs(t, err, errNotAuthorized)
   612  		require.NotZero(t, backoff)
   613  	})
   614  
   615  	t.Run("transport closing - not connected", func(t *testing.T) {
   616  		backoff := testBackoff(1)
   617  		done, err := handleError(ctx, logger, &backoff, false, errTransportClosing)
   618  		require.False(t, done)
   619  		require.NoError(t, err)
   620  		require.NotZero(t, backoff)
   621  	})
   622  
   623  	t.Run("transport closing - not connected - cancelled context", func(t *testing.T) {
   624  		ctx, cancel := context.WithCancel(ctx)
   625  		cancel()
   626  
   627  		backoff := testBackoff(time.Second)
   628  		done, err := handleError(ctx, logger, &backoff, false, errTransportClosing)
   629  		require.True(t, done)
   630  		require.ErrorIs(t, err, context.Canceled)
   631  		require.NotZero(t, backoff)
   632  	})
   633  
   634  	t.Run("transport closing - connected", func(t *testing.T) {
   635  		backoff := testBackoff(1)
   636  		done, err := handleError(ctx, logger, &backoff, true, errTransportClosing)
   637  		require.False(t, done)
   638  		require.NoError(t, err)
   639  		require.Zero(t, backoff)
   640  	})
   641  
   642  	t.Run("probe unregistered - not connected", func(t *testing.T) {
   643  		backoff := testBackoff(1)
   644  		done, err := handleError(ctx, logger, &backoff, false, errProbeUnregistered)
   645  		require.False(t, done)
   646  		require.NoError(t, err)
   647  		require.NotZero(t, backoff)
   648  	})
   649  
   650  	t.Run("probe unregistered - connected", func(t *testing.T) {
   651  		backoff := testBackoff(1)
   652  		done, err := handleError(ctx, logger, &backoff, true, errProbeUnregistered)
   653  		require.False(t, done)
   654  		require.NoError(t, err)
   655  		require.Zero(t, backoff)
   656  	})
   657  }
   658  
   659  func TestProbeTenantCh(t *testing.T) {
   660  	t.Run("sends probe tenant ID after registration", func(t *testing.T) {
   661  		probeTenantCh := make(chan *sm.Probe, 1)
   662  		u, err := NewUpdater(UpdaterOptions{
   663  			Conn:           new(grpc.ClientConn),
   664  			PromRegisterer: prometheus.NewPedanticRegistry(),
   665  			Publisher:      channelPublisher(make(chan pusher.Payload)),
   666  			TenantCh:       make(chan<- sm.Tenant),
   667  			ProbeCh:        probeTenantCh,
   668  			Logger:         testhelper.Logger(t),
   669  		})
   670  
   671  		require.NoError(t, err)
   672  
   673  		u.probe = &sm.Probe{Id: 1, TenantId: 42}
   674  		u.notifyProbeTenant()
   675  
   676  		got := <-probeTenantCh
   677  		require.Equal(t, u.probe.TenantId, got.TenantId)
   678  	})
   679  
   680  	t.Run("only sends once across multiple registrations", func(t *testing.T) {
   681  		probeCh := make(chan *sm.Probe, 1)
   682  		u, err := NewUpdater(UpdaterOptions{
   683  			Conn:           new(grpc.ClientConn),
   684  			PromRegisterer: prometheus.NewPedanticRegistry(),
   685  			Publisher:      channelPublisher(make(chan pusher.Payload)),
   686  			TenantCh:       make(chan<- sm.Tenant),
   687  			ProbeCh:        probeCh,
   688  			Logger:         testhelper.Logger(t),
   689  		})
   690  
   691  		require.NoError(t, err)
   692  
   693  		u.probe = &sm.Probe{Id: 1, TenantId: 42}
   694  		u.notifyProbeTenant()
   695  
   696  		u.probe = &sm.Probe{Id: 2, TenantId: 43}
   697  		u.notifyProbeTenant()
   698  
   699  		// Confirm that the returned value is the first probe id sent to the channel
   700  		got := <-probeCh
   701  		require.Equal(t, got.TenantId, int64(42))
   702  
   703  		// Confirm no second send
   704  		_, ok := <-probeCh
   705  		require.False(t, ok)
   706  	})
   707  
   708  	t.Run("nil channel is safe when feature disabled", func(t *testing.T) {
   709  		u, err := NewUpdater(UpdaterOptions{
   710  			Conn:           new(grpc.ClientConn),
   711  			PromRegisterer: prometheus.NewPedanticRegistry(),
   712  			Publisher:      channelPublisher(make(chan pusher.Payload)),
   713  			TenantCh:       make(chan<- sm.Tenant),
   714  			Logger:         testhelper.Logger(t),
   715  		})
   716  
   717  		require.NoError(t, err)
   718  
   719  		u.probe = &sm.Probe{Id: 1, TenantId: 42}
   720  
   721  		require.NotPanics(t, func() {
   722  			u.notifyProbeTenant()
   723  		})
   724  	})
   725  }