Repository: prometheus/prometheus
File: notifier/manager_test.go (Go)

Code before:

func TestStop_DrainingEnabled(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		const alertmanagerURL = "http://alertmanager:9093/api/v2/alerts"

		alertsReceived := atomic.NewInt64(0)
		handlerStarted := make(chan struct{}, 1)

		// Fake Do function that simulates alertmanager responding slowly but successfully.
		// This paces request handling so we can queue alert-2 while alert-1 is still in-flight.
		fakeDo := func(_ context.Context, _ *http.Client, req *http.Request) (*http.Response, error) {
			var alerts []*Alert
			b, err := io.ReadAll(req.Body)
			if err != nil {
				return nil, fmt.Errorf("read request body: %w", err)
			}
			if err := json.Unmarshal(b, &alerts); err != nil {
				return nil, fmt.Errorf("unmarshal request body: %w", err)
			}

			// Signal arrival.
			handlerStarted <- struct{}{}

			alertsReceived.Add(int64(len(alerts)))

			// Respond slowly to allow time for alert-2 to be queued while this request is in-flight.
			time.Sleep(100 * time.Millisecond)

			return &http.Response{
				StatusCode: http.StatusOK,
				Body:       io.NopCloser(bytes.NewBuffer(nil)),
			}, nil
		}

		reg := prometheus.NewRegistry()
		m := NewManager(
			&Options{
				QueueCapacity:   10,
				DrainOnShutdown: true,
				Registerer:      reg,
				Do:              fakeDo,
			},
			model.UTF8Validation,
			nil,
		)

		m.alertmanagers = make(map[string]*alertmanagerSet)

		am1Cfg := config.DefaultAlertmanagerConfig
		am1Cfg.Timeout = model.Duration(time.Second)
		m.alertmanagers["1"] = newTestAlertmanagerSet(&am1Cfg, nil, m.opts, m.metrics, alertmanagerURL)

Review comment (performance):
"Slowly" is interesting word given in synctest time sleep blocks goroutines. Should we just say `Block on time so  alert-2 is queued while this request is in-flight.`

Code after:
			b, err := io.ReadAll(req.Body)
			if err != nil {
				return nil, fmt.Errorf("read request body: %w", err)
			}
			if err := json.Unmarshal(b, &alerts); err != nil {
				return nil, fmt.Errorf("unmarshal request body: %w", err)
			}
			alertsReceived.Add(int64(len(alerts)))

			// Signal arrival.
			handlerStarted <- struct{}{}

			// Block to allow for alert-2 to be queued while this request is in-flight.
			time.Sleep(100 * time.Millisecond)

			return &http.Response{
				StatusCode: http.StatusOK,
				Body:       io.NopCloser(bytes.NewBuffer(nil)),
			}, nil
		}

		reg := prometheus.NewRegistry()
		m := NewManager(
			&Options{
				QueueCapacity:   10,
				DrainOnShutdown: true,
				Registerer:      reg,
				Do:              fakeDo,
			},
			model.UTF8Validation,
			nil,
		)

		m.alertmanagers = make(map[string]*alertmanagerSet)

		am1Cfg := config.DefaultAlertmanagerConfig
		am1Cfg.Timeout = model.Duration(time.Second)
		m.alertmanagers["1"] = newTestAlertmanagerSet(&am1Cfg, nil, m.opts, m.metrics, alertmanagerURL)

		for _, ams := range m.alertmanagers {
			ams.startSendLoops(ams.ams)
		}

		go m.Run(nil)

		// Queue two alerts. The first should be immediately sent to the receiver.
		m.Send(&Alert{Labels: labels.FromStrings(labels.AlertName, "alert-1")})

		// Wait for receiver to get the first request.
		<-handlerStarted