o11y/pkg/query-service/rules/manager_test.go

package rules

import (
	"context"
	"encoding/json"
	"math"
	"strconv"
	"testing"
	"time"

	"github.com/hanzoai/o11y/pkg/alertmanager"
	"github.com/hanzoai/o11y/pkg/alertmanager/alertmanagerserver"
	alertmanagermock "github.com/hanzoai/o11y/pkg/alertmanager/alertmanagertest"
	"github.com/hanzoai/o11y/pkg/instrumentation/instrumentationtest"
	"github.com/hanzoai/o11y/pkg/prometheus"
	"github.com/hanzoai/o11y/pkg/prometheus/prometheustest"
	"github.com/hanzoai/o11y/pkg/sqlstore"
	"github.com/hanzoai/o11y/pkg/sqlstore/sqlstoretest"
	"github.com/hanzoai/o11y/pkg/telemetrystore"
	"github.com/hanzoai/o11y/pkg/telemetrystore/telemetrystoretest"
	"github.com/hanzoai/o11y/pkg/types/alertmanagertypes"
	"github.com/hanzoai/o11y/pkg/types/metrictypes"
	"github.com/hanzoai/o11y/pkg/valuer"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	dsmock "github.com/hanzo-ds/mock"
)

func TestManager_TestNotification_SendUnmatched_ThresholdRule(t *testing.T) {
	target := 10.0
	recovery := 5.0

	for _, tc := range TcTestNotiSendUnmatchedThresholdRule {
		t.Run(tc.Name, func(t *testing.T) {
			rule := ThresholdRuleAtLeastOnceValueAbove(target, &recovery)

			// Marshal rule to JSON as TestNotification expects
			ruleBytes, err := json.Marshal(rule)
			require.NoError(t, err)

			orgID := valuer.GenerateUUID()

			// for saving temp alerts that are triggered via TestNotification
			triggeredTestAlerts := []map[*alertmanagertypes.PostableAlert][]string{}

			// Create manager using test factory with hooks
			mgr := NewTestManager(t, &TestManagerOptions{
				AlertmanagerHook: func(am alertmanager.Alertmanager) {
					mockAM := am.(*alertmanagermock.MockAlertmanager)
					// mock set notification config
					mockAM.On("SetNotificationConfig", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
					mockAM.On("Config").Return(alertmanagerserver.Config{ExternalURL: mustParseURL(t, "http://localhost:8080")})
					// for saving temp alerts that are triggered via TestNotification
					if tc.ExpectAlerts > 0 {
						mockAM.On("TestAlert", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
							triggeredTestAlerts = append(triggeredTestAlerts, args.Get(3).(map[*alertmanagertypes.PostableAlert][]string))
						}).Return(nil).Times(tc.ExpectAlerts)
					}
				},
				SqlStoreHook: func(store sqlstore.SQLStore) {
					mockStore := store.(*sqlstoretest.Provider)
					// Mock the organizations query that SendAlerts makes
					// Bun generates: SELECT id FROM organizations LIMIT 1 (or SELECT "id" FROM "organizations" LIMIT 1)
					orgRows := mockStore.Mock().NewRows([]string{"id"}).AddRow(orgID.StringValue())
					// Match bun's generated query pattern - bun may quote identifiers
					mockStore.Mock().ExpectQuery("SELECT (.+) FROM (.+)organizations(.+) LIMIT (.+)").WillReturnRows(orgRows)
				},
				TelemetryStoreHook: func(store telemetrystore.TelemetryStore) {
					mockStore := store.(*telemetrystoretest.Provider)
					// Set up mock data for telemetry store
					cols := make([]dsmock.ColumnType, 0)
					cols = append(cols, dsmock.ColumnType{Name: "value", Type: "Float64"})
					cols = append(cols, dsmock.ColumnType{Name: "attr", Type: "String"})
					cols = append(cols, dsmock.ColumnType{Name: "ts", Type: "DateTime"})

					alertDataRows := dsmock.NewRows(cols, tc.Values)

					mock := mockStore.Mock()
					// Mock metadata queries for FetchTemporalityAndTypeMulti
					// First query: fetchMetricsTemporalityAndType (from o11y_metrics time series table)
					metadataCols := []dsmock.ColumnType{
						{Name: "metric_name", Type: "String"},
						{Name: "temporality", Type: "String"},
						{Name: "type", Type: "String"},
						{Name: "is_monotonic", Type: "Bool"},
					}
					metadataRows := dsmock.NewRows(metadataCols, [][]any{
						{"probe_success", metrictypes.Unspecified, metrictypes.GaugeType, false},
					})
					mock.ExpectQuery("*distributed_time_series_v4*").WithArgs(nil, nil, nil).WillReturnRows(metadataRows)
					// Second query: fetchMeterSourceMetricsTemporalityAndType (from o11y_meter table)
					emptyMetadataRows := dsmock.NewRows(metadataCols, [][]any{})
					mock.ExpectQuery("*meter*").WithArgs(nil).WillReturnRows(emptyMetadataRows)

					// Generate query arguments for the metric query
					evalTime := time.Now().UTC()
					evalWindow := 5 * time.Minute
					evalDelay := time.Duration(0)
					queryArgs := GenerateMetricQueryCHArgs(
						evalTime,
						evalWindow,
						evalDelay,
						"probe_success",
						metrictypes.Unspecified,
					)

					mock.ExpectQuery("*WITH __temporal_aggregation_cte*").
						WithArgs(queryArgs...).
						WillReturnRows(alertDataRows)
				},
			})

			count, err := mgr.TestNotification(context.Background(), orgID, string(ruleBytes))
			require.Nil(t, err)
			assert.Equal(t, tc.ExpectAlerts, count)

			if tc.ExpectAlerts > 0 {
				// check if the alert has been triggered
				require.Len(t, triggeredTestAlerts, 1)
				var gotAlerts []*alertmanagertypes.PostableAlert
				for a := range triggeredTestAlerts[0] {
					gotAlerts = append(gotAlerts, a)
				}
				require.Len(t, gotAlerts, tc.ExpectAlerts)
				// check if the alert has triggered with correct threshold value
				if tc.ExpectValue != 0 {
					assert.Equal(t, strconv.FormatFloat(tc.ExpectValue, 'f', -1, 64), gotAlerts[0].Annotations["value"])
				}
			} else {
				// check if no alerts have been triggered
				assert.Empty(t, triggeredTestAlerts)
			}
		})
	}
}

func TestManager_TestNotification_SendUnmatched_PromRule(t *testing.T) {
	target := 10.0

	for _, tc := range TcTestNotificationSendUnmatchedPromRule {
		t.Run(tc.Name, func(t *testing.T) {
			// Capture base time once per test case to ensure consistent timestamps
			baseTime := time.Now().UTC()

			rule := BuildPromAtLeastOnceValueAbove(target, nil)

			// Marshal rule to JSON as TestNotification expects
			ruleBytes, err := json.Marshal(rule)
			require.NoError(t, err)

			orgID := valuer.GenerateUUID()

			// for saving temp alerts that are triggered via TestNotification
			triggeredTestAlerts := []map[*alertmanagertypes.PostableAlert][]string{}

			// Variable to store promProvider for cleanup
			var promProvider *prometheustest.Provider

			// Create manager using test factory with hooks
			mgr := NewTestManager(t, &TestManagerOptions{
				AlertmanagerHook: func(am alertmanager.Alertmanager) {
					mockAM := am.(*alertmanagermock.MockAlertmanager)
					// mock set notification config
					mockAM.On("SetNotificationConfig", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
					mockAM.On("Config").Return(alertmanagerserver.Config{ExternalURL: mustParseURL(t, "http://localhost:8080")})
					// for saving temp alerts that are triggered via TestNotification
					if tc.ExpectAlerts > 0 {
						mockAM.On("TestAlert", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
							triggeredTestAlerts = append(triggeredTestAlerts, args.Get(3).(map[*alertmanagertypes.PostableAlert][]string))
						}).Return(nil).Times(tc.ExpectAlerts)
					}
				},
				SqlStoreHook: func(store sqlstore.SQLStore) {
					mockStore := store.(*sqlstoretest.Provider)
					// Mock the organizations query that SendAlerts makes
					orgRows := mockStore.Mock().NewRows([]string{"id"}).AddRow(orgID.StringValue())
					mockStore.Mock().ExpectQuery("SELECT (.+) FROM (.+)organizations(.+) LIMIT (.+)").WillReturnRows(orgRows)
				},
				TelemetryStoreHook: func(store telemetrystore.TelemetryStore) {
					mockStore := store.(*telemetrystoretest.Provider)

					// Set up Prometheus-specific mock data
					// Fingerprint columns for Prometheus queries
					fingerprintCols := []dsmock.ColumnType{
						{Name: "fingerprint", Type: "UInt64"},
						{Name: "any(labels)", Type: "String"},
					}

					// Samples columns for Prometheus queries
					samplesCols := []dsmock.ColumnType{
						{Name: "metric_name", Type: "String"},
						{Name: "fingerprint", Type: "UInt64"},
						{Name: "unix_milli", Type: "Int64"},
						{Name: "value", Type: "Float64"},
						{Name: "flags", Type: "UInt32"},
					}

					// Calculate query time range similar to Prometheus rule tests
					// TestNotification uses time.Now().UTC() for evaluation
					// We calculate the query window based on current time to match what the actual evaluation will use
					evalTime := baseTime
					evalWindowMs := int64(5 * 60 * 1000) // 5 minutes in ms
					evalTimeMs := evalTime.UnixMilli()
					queryStart := ((evalTimeMs-2*evalWindowMs)/60000)*60000 + 1 // truncate to minute + 1ms
					queryEnd := (evalTimeMs / 60000) * 60000                    // truncate to minute

					// Create fingerprint data
					fingerprint := uint64(12345)
					labelsJSON := `{"__name__":"test_metric"}`
					fingerprintData := [][]any{
						{fingerprint, labelsJSON},
					}
					fingerprintRows := dsmock.NewRows(fingerprintCols, fingerprintData)

					// Create samples data from test case values, calculating timestamps relative to baseTime
					validSamplesData := make([][]any, 0)
					for _, v := range tc.Values {
						// Skip NaN and Inf values in the samples data
						if math.IsNaN(v.Value) || math.IsInf(v.Value, 0) {
							continue
						}
						// Calculate timestamp relative to baseTime
						sampleTimestamp := baseTime.Add(v.Offset).UnixMilli()
						validSamplesData = append(validSamplesData, []any{
							"test_metric",
							fingerprint,
							sampleTimestamp,
							v.Value,
							uint32(0), // flags - 0 means normal value
						})
					}
					samplesRows := dsmock.NewRows(samplesCols, validSamplesData)

					mock := mockStore.Mock()

					// Mock the fingerprint query (for Prometheus label matching)
					mock.ExpectQuery("SELECT fingerprint, any").
						WithArgs("test_metric", "__name__", "test_metric").
						WillReturnRows(fingerprintRows)

					// Mock the samples query (for Prometheus metric data)
					mock.ExpectQuery("SELECT metric_name, fingerprint, unix_milli").
						WithArgs(
							"test_metric",
							"test_metric",
							"__name__",
							"test_metric",
							queryStart,
							queryEnd,
						).
						WillReturnRows(samplesRows)

					// Create Prometheus provider for this test
					promProvider = prometheustest.New(context.Background(), instrumentationtest.New().ToProviderSettings(), prometheus.Config{Timeout: 2 * time.Minute}, store)
				},
				ManagerOptionsHook: func(opts *ManagerOptions) {
					// Set Prometheus provider for PromQL queries
					if promProvider != nil {
						opts.Prometheus = promProvider
					}
				},
			})

			count, err := mgr.TestNotification(context.Background(), orgID, string(ruleBytes))
			require.Nil(t, err)
			assert.Equal(t, tc.ExpectAlerts, count)

			if tc.ExpectAlerts > 0 {
				// check if the alert has been triggered
				require.Len(t, triggeredTestAlerts, 1)
				var gotAlerts []*alertmanagertypes.PostableAlert
				for a := range triggeredTestAlerts[0] {
					gotAlerts = append(gotAlerts, a)
				}
				require.Len(t, gotAlerts, tc.ExpectAlerts)
				// check if the alert has triggered with correct threshold value
				if tc.ExpectValue != 0 && !math.IsNaN(tc.ExpectValue) && !math.IsInf(tc.ExpectValue, 0) {
					assert.Equal(t, strconv.FormatFloat(tc.ExpectValue, 'f', -1, 64), gotAlerts[0].Annotations["value"])
				}
			} else {
				// check if no alerts have been triggered
				assert.Empty(t, triggeredTestAlerts)
			}

			promProvider.Close()
		})
	}
}