Alerting: Extend recording rules test to exercise writing with data sources. (#101775)

The change to use WriteDatasource was done in a previous commit, this adds a
test case using DatasourceWriter, in addition to the one using PrometheusWriter.

pkg/services/ngalert/schedule/recording_rule_test.go

@@ -20,6 +20,8 @@ import (
 
 	"github.com/grafana/grafana/pkg/expr"
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	dsfakes "github.com/grafana/grafana/pkg/services/datasources/fakes"
 	"github.com/grafana/grafana/pkg/services/ngalert/metrics"
 	models "github.com/grafana/grafana/pkg/services/ngalert/models"
 	"github.com/grafana/grafana/pkg/services/ngalert/writer"
@@ -173,18 +175,33 @@ func blankRecordingRuleForTests(ctx context.Context) *recordingRule {
 }
 
 func TestRecordingRule_Integration(t *testing.T) {
+	t.Run("with prometheus writer", func(t *testing.T) {
+		writeTarget := writer.NewTestRemoteWriteTarget(t)
+		defer writeTarget.Close()
+		writerReg := prometheus.NewPedanticRegistry()
+		writer := setupPrometheusWriter(t, writeTarget, writerReg)
+		testRecordingRule_Integration(t, writeTarget, writer, writerReg, "")
+	})
+	t.Run("with datasource writer", func(t *testing.T) {
+		writeTarget := writer.NewTestRemoteWriteTarget(t)
+		defer writeTarget.Close()
+		writerReg := prometheus.NewPedanticRegistry()
+		writer := setupDatasourceWriter(t, writeTarget, writerReg, "ds-uid")
+		testRecordingRule_Integration(t, writeTarget, writer, writerReg, "ds-uid")
+	})
+}
+
+func testRecordingRule_Integration(t *testing.T, writeTarget *writer.TestRemoteWriteTarget, writer RecordingWriter, writerReg *prometheus.Registry, dsUID string) {
 	gen := models.RuleGen.With(models.RuleGen.WithAllRecordingRules(), models.RuleGen.WithOrgID(123))
 	ruleStore := newFakeRulesStore()
 	reg := prometheus.NewPedanticRegistry()
 	sch := setupScheduler(t, ruleStore, nil, reg, nil, nil, nil)
-	writeTarget := writer.NewTestRemoteWriteTarget(t)
+	sch.recordingWriter = writer
-	defer writeTarget.Close()
-	writerReg := prometheus.NewPedanticRegistry()
-	sch.recordingWriter = setupWriter(t, writeTarget, writerReg)
 
 	t.Run("rule that succeeds", func(t *testing.T) {
 		writeTarget.Reset()
 		rule := gen.With(withQueryForHealth("ok")).GenerateRef()
+		rule.Record.TargetDatasourceUID = dsUID
 		ruleStore.PutRule(context.Background(), rule)
 		folderTitle := ruleStore.getNamespaceTitle(rule.NamespaceUID)
 		ruleFactory := ruleFactoryFromScheduler(sch)
@@ -543,7 +560,7 @@ func withQueryForHealth(health string) models.AlertRuleMutator {
 	}
 }
 
-func setupWriter(t *testing.T, target *writer.TestRemoteWriteTarget, reg prometheus.Registerer) *writer.PrometheusWriter {
+func setupPrometheusWriter(t *testing.T, target *writer.TestRemoteWriteTarget, reg prometheus.Registerer) *writer.PrometheusWriter {
 	provider := testClientProvider{}
 	m := metrics.NewNGAlert(reg)
 	wr, err := writer.NewPrometheusWriterWithSettings(target.ClientSettings(), provider, clock.NewMock(), log.NewNopLogger(), m.GetRemoteWriterMetrics())
@@ -551,6 +568,27 @@ func setupWriter(t *testing.T, target *writer.TestRemoteWriteTarget, reg prometh
 	return wr
 }
 
+func setupDatasourceWriter(t *testing.T, target *writer.TestRemoteWriteTarget, reg prometheus.Registerer, dsUID string) *writer.DatasourceWriter {
+	provider := testClientProvider{}
+	m := metrics.NewNGAlert(reg)
+
+	dss := &dsfakes.FakeDataSourceService{}
+	p1, _ := dss.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+		UID:  dsUID,
+		Type: datasources.DS_PROMETHEUS,
+	})
+	p1.URL = target.DatasourceURL()
+
+	cfg := writer.DatasourceWriterConfig{
+		Timeout:               time.Second * 5,
+		DefaultDatasourceUID:  "",
+		RemoteWritePathSuffix: writer.RemoteWriteSuffix,
+	}
+
+	return writer.NewDatasourceWriter(cfg, dss, provider, clock.NewMock(),
+		log.New("test"), m.GetRemoteWriterMetrics())
+}
+
 type testClientProvider struct{}
 
 func (t testClientProvider) New(options ...httpclient.Options) (*http.Client, error) {

pkg/services/ngalert/schedule/schedule.go

@@ -49,7 +49,6 @@ type RulesStore interface {
 }
 
 type RecordingWriter interface {
-	Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error
 	WriteDatasource(ctx context.Context, dsUID string, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error
 }
 

pkg/services/ngalert/writer/fake.go

@@ -12,14 +12,6 @@ type FakeWriter struct {
 	WriteFunc func(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error
 }
 
-func (w FakeWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error {
-	if w.WriteFunc == nil {
-		return nil
-	}
-
-	return w.WriteFunc(ctx, name, t, frames, orgID, extraLabels)
-}
-
 func (w FakeWriter) WriteDatasource(ctx context.Context, dsUID string, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error {
 	if w.WriteFunc == nil {
 		return nil

pkg/services/ngalert/writer/testing.go

@@ -12,7 +12,10 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const RemoteWriteEndpoint = "/api/v1/write"
+const RemoteWritePrefix = "/api/v1"
+const RemoteWriteSuffix = "/write"
+
+const RemoteWriteEndpoint = RemoteWritePrefix + RemoteWriteSuffix
 
 type TestRemoteWriteTarget struct {
 	srv *httptest.Server
@@ -59,6 +62,10 @@ func (s *TestRemoteWriteTarget) Close() {
 	s.srv.Close()
 }
 
+func (s *TestRemoteWriteTarget) DatasourceURL() string {
+	return s.srv.URL + RemoteWritePrefix
+}
+
 func (s *TestRemoteWriteTarget) ClientSettings() setting.RecordingRuleSettings {
 	return setting.RecordingRuleSettings{
 		URL:               s.srv.URL + RemoteWriteEndpoint,