worker/app/models/check_aggregator_test.go

package models_test

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/datatypes"
	"gorm.io/gorm"

	"rocketgit.ru/rsmon/worker/app/models"
	"rocketgit.ru/rsmon/worker/internal/wire"
)

// seedAggregatorWorld creates the minimum fixture the aggregator tests
// need: a plan, account, group, monitor, and an http check whose
// RequireQuorum / AggregationWindowSeconds are set per call. The check
// is created with state=UNK so the test can observe the aggregator's
// effect on Check.State directly.
//
// Returns the freshly-created monitor + check; the check is what every
// test below mutates (RequireQuorum, AggregationWindowSeconds) and then
// asserts on. Cleanup is the caller's responsibility — most tests call
// models.Drop() at the top instead.
func seedAggregatorWorld(t *testing.T, quorum, windowSeconds int) (models.Monitor, models.Check) {
	t.Helper()

	plan := models.Plan{Name: "agg-plan", Default: true}
	require.NoError(t, models.DB().Create(&plan).Error)
	acc := models.Account{Name: "agg-acc", PlanID: &plan.ID}
	require.NoError(t, models.DB().Create(&acc).Error)
	grp := &models.Group{AccountID: acc.ID, Name: "agg"}
	require.NoError(t, models.DB().Create(grp).Error)

	mon := models.Monitor{
		Name:    stringPtrAgg("agg.test"),
		Host:    "agg.test",
		GroupID: grp.ID,
		Enabled: true,
	}
	require.NoError(t, models.DB().Create(&mon).Error)

	enTrue := true
	check := models.Check{
		MonitorID:                mon.ID,
		Kind:                     "http",
		Interval:                 60,
		Enabled:                  &enTrue,
		State:                    "UNK",
		Settings:                 datatypes.JSON([]byte(`{}`)),
		RequireQuorum:            quorum,
		AggregationWindowSeconds: windowSeconds,
	}
	require.NoError(t, models.DB().Create(&check).Error)
	return mon, check
}

func stringPtrAgg(s string) *string { return &s }

// makeReport constructs a wire.CheckResultReport with sensible defaults
// for the OK or not-OK case. Tests use this to push results through
// ApplyRemoteCheckResult exactly the way a real worker would.
func makeReport(checkID, monitorID int64, state string) wire.CheckResultReport {
	return wire.CheckResultReport{
		JobID:     "job-" + state,
		CheckID:   checkID,
		MonitorID: monitorID,
		State:     state,
	}
}

// regionResultWithErr is regionResultFor plus an error message. Used
// when the test wants to verify that the aggregator forwards the row's
// error string onto Check.Error (mimics a real worker reporting
// state=ERR with a diagnostic message).
func regionResultWithErr(t *testing.T, checkID int64, region, state, errMsg string, age time.Duration) models.CheckRegionResult {
	t.Helper()
	row := regionResultFor(t, checkID, region, state, age)
	require.NoError(t, models.DB().Model(&row).UpdateColumn("error", errMsg).Error)
	return row
}

// regionResultFor inserts a single CheckRegionResult row whose
// created_at and executed_at are both backdated by `age`, so the
// aggregator's window-based watermark picks it up immediately without
// needing a real time.Sleep. Returned row has its DB-assigned ID
// populated.
func regionResultFor(t *testing.T, checkID int64, region string, state string, age time.Duration) models.CheckRegionResult {
	t.Helper()
	// RegionCode has a FK to regions.code, so the region must exist
	// before the result row is inserted. seedRegion is idempotent.
	seedRegion(t, region)
	row := models.CheckRegionResult{
		CheckID:    checkID,
		RegionCode: region,
		ExecutedAt: time.Now().Add(-age),
		State:      state,
	}
	require.NoError(t, models.DB().Create(&row).Error)
	// Backdate CreatedAt too — the aggregator SQL keys on
	// check_region_results.created_at (see CheckAggregatorTick). GORM
	// auto-sets CreatedAt on insert, so we have to UPDATE it post-hoc.
	require.NoError(t, models.DB().Model(&row).UpdateColumns(map[string]interface{}{
		"created_at": time.Now().Add(-age),
		"updated_at": time.Now().Add(-age),
	}).Error)
	return row
}

// loadCheck re-reads a Check row by ID — used after the aggregator
// runs so the test asserts against the post-tick state.
func loadCheck(t *testing.T, id int64) models.Check {
	t.Helper()
	var c models.Check
	require.NoError(t, models.DB().First(&c, id).Error)
	return c
}

// countPendingResults returns how many CheckRegionResult rows for
// checkID have aggregated_at IS NULL — the working set the next
// aggregator tick would consider.
func countPendingResults(t *testing.T, checkID int64) int64 {
	t.Helper()
	var n int64
	require.NoError(t, models.DB().Model(&models.CheckRegionResult{}).
		Where("check_id = ? AND aggregated_at IS NULL", checkID).
		Count(&n).Error)
	return n
}

// ---------------------------------------------------------------------------
// ApplyRemoteCheckResult: regression tests for the QuorumEnabled split.
// ---------------------------------------------------------------------------

// TestApplyRemoteCheckResult_DirectWhenQuorumOne pins the legacy path:
// when RequireQuorum==1 (the default), ApplyRemoteCheckResult still
// writes Check.State synchronously, exactly the way it did before Phase
// 3. This is the regression guard for the in-process RKN scheduler
// tests in internal/rknscheduler.
func TestApplyRemoteCheckResult_DirectWhenQuorumOne(t *testing.T) {
	models.Drop()
	models.Migrate()

	mon, check := seedAggregatorWorld(t, 1, 5)
	// Region must exist because CheckRegionResult has a FK to
	// regions.code (seeded by Migrate, but the test region is custom).
	seedRegion(t, "ru-msk")

	require.NoError(t, models.ApplyRemoteCheckResult(
		makeReport(check.ID, mon.ID, "OK"),
		"ru-msk",
	))

	got := loadCheck(t, check.ID)
	assert.Equal(t, "OK", got.State, "quorum=1 must keep the direct State update")
	assert.NotNil(t, got.LastEnd, "legacy path must keep stamping last_end")
	// One region result inserted with aggregated_at=NULL.
	assert.EqualValues(t, 1, countPendingResults(t, check.ID),
		"the region result row is always inserted even on the legacy path")
}

func TestApplyRemoteCheckResultFromWorkerDeduplicatesFailureReplay(t *testing.T) {
	models.Drop()
	models.Migrate()

	mon, check := seedAggregatorWorld(t, 1, 5)
	seedRegion(t, "ru-msk")
	worker := models.WorkerNode{
		WorkerID:   "replay-worker",
		RegionCode: "ru-msk",
		AuthToken:  "replay-token",
	}
	require.NoError(t, models.DB().Create(&worker).Error)
	report := makeReport(check.ID, mon.ID, "ERR")
	report.JobID = "same-failed-execution"

	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))
	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))

	got := loadCheck(t, check.ID)
	assert.Equal(t, 1, got.Fails)
	assert.EqualValues(t, 1, countPendingResults(t, check.ID))
	var attempts int64
	require.NoError(t, models.DB().Model(&models.CheckAttempt{}).Where("job_id = ?", report.JobID).Count(&attempts).Error)
	assert.EqualValues(t, 1, attempts)
}

func TestApplyRemoteCheckResultFromWorkerDeduplicatesQuorumReplay(t *testing.T) {
	models.Drop()
	models.Migrate()

	mon, check := seedAggregatorWorld(t, 2, 5)
	seedRegion(t, "ru-msk")
	worker := models.WorkerNode{
		WorkerID:   "quorum-replay-worker",
		RegionCode: "ru-msk",
		AuthToken:  "quorum-replay-token",
	}
	require.NoError(t, models.DB().Create(&worker).Error)
	report := makeReport(check.ID, mon.ID, "OK")
	report.JobID = "same-quorum-execution"

	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))
	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))

	got := loadCheck(t, check.ID)
	assert.Equal(t, "UNK", got.State, "quorum state remains owned by the aggregator")
	assert.EqualValues(t, 1, countPendingResults(t, check.ID))
	var attempts int64
	require.NoError(t, models.DB().Model(&models.CheckAttempt{}).Where("job_id = ?", report.JobID).Count(&attempts).Error)
	assert.EqualValues(t, 1, attempts)
}

func TestApplyRemoteCheckResultFromWorkerDeduplicatesFailReplay(t *testing.T) {
	models.Drop()
	models.Migrate()

	mon, check := seedAggregatorWorld(t, 1, 5)
	seedRegion(t, "ru-msk")
	worker := models.WorkerNode{
		WorkerID:   "fail-replay-worker",
		RegionCode: "ru-msk",
		AuthToken:  "fail-replay-token",
	}
	require.NoError(t, models.DB().Create(&worker).Error)
	report := makeReport(check.ID, mon.ID, "FAIL")
	report.JobID = "same-fail-execution"

	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))
	require.NoError(t, models.ApplyRemoteCheckResultFromWorker(report, worker.RegionCode, &worker))

	got := loadCheck(t, check.ID)
	assert.Equal(t, "FAIL", got.State)
	assert.Equal(t, 0, got.Fails, "FAIL resets consecutive ERR tracking")
	assert.EqualValues(t, 1, countPendingResults(t, check.ID))
	var attempts int64
	require.NoError(t, models.DB().Model(&models.CheckAttempt{}).Where("job_id = ?", report.JobID).Count(&attempts).Error)
	assert.EqualValues(t, 1, attempts)
}

// TestApplyRemoteCheckResult_BuffersWhenQuorumN pins the new path:
// when RequireQuorum > 1, ApplyRemoteCheckResult does NOT touch
// Check.State — it only inserts the CheckRegionResult row. The check
// stays at its initial UNK and the unaggregated row count grows by
// exactly 1 per call.
func TestApplyRemoteCheckResult_BuffersWhenQuorumN(t *testing.T) {
	models.Drop()
	models.Migrate()

	mon, check := seedAggregatorWorld(t, 3, 5)
	seedRegion(t, "ru-msk")
	seedRegion(t, "us-east")

	require.NoError(t, models.ApplyRemoteCheckResult(
		makeReport(check.ID, mon.ID, "OK"), "ru-msk",
	))
	require.NoError(t, models.ApplyRemoteCheckResult(
		makeReport(check.ID, mon.ID, "ERR"), "us-east",
	))

	got := loadCheck(t, check.ID)
	assert.Equal(t, "UNK", got.State,
		"quorum>1 must NOT touch Check.State — the aggregator owns it")
	assert.EqualValues(t, 2, countPendingResults(t, check.ID),
		"two results buffered, both with aggregated_at=NULL")
}

// ---------------------------------------------------------------------------
// CheckAggregatorTick: rule tests.
// ---------------------------------------------------------------------------

// TestAggregator_QuorumOK: with RequireQuorum=2 and two OK results
// buffered, the aggregator must decide OK and stamp AggregatedAt on
// both contributing rows.
func TestAggregator_QuorumOK(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 1)
	regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)
	regionResultFor(t, check.ID, "us-east", "OK", 4*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 1, n, "one check aggregated this tick")

	got := loadCheck(t, check.ID)
	assert.Equal(t, "OK", got.State, "2 OK results >= quorum=2 → state=OK")
	assert.NotNil(t, got.LastOk, "OK decision must stamp last_ok")
	assert.EqualValues(t, 0, got.Fails, "fails must reset on OK")
	assert.EqualValues(t, 0, countPendingResults(t, check.ID),
		"both contributing rows must be stamped aggregated_at")
}

// TestAggregator_QuorumFail: with RequireQuorum=2 and two ERR results
// buffered, the aggregator must decide ERR and surface the most recent
// row's error message on Check.Error.
func TestAggregator_QuorumFail(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 1)
	regionResultWithErr(t, check.ID, "ru-msk", "ERR", "connection refused", 5*time.Second)
	regionResultWithErr(t, check.ID, "us-east", "ERR", "timeout", 4*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 1, n)

	got := loadCheck(t, check.ID)
	assert.Equal(t, "ERR", got.State, "2 ERR >= quorum=2 → state=ERR")
	assert.NotNil(t, got.LastFail)
	assert.NotNil(t, got.Error, "ERR decision must carry an error message from the rows")
	assert.Contains(t, *got.Error, "timeout",
		"aggregator should surface the latest row's error message")
	assert.EqualValues(t, 0, countPendingResults(t, check.ID))
}

// TestAggregator_DegradedWhenPartial: with RequireQuorum=3 and 1 OK +
// 2 ERR (mixed within window), neither side reaches the quorum of 3 so
// the aggregator must decide DEGRADED. The state must NOT silently
// become OK or ERR.
func TestAggregator_DegradedWhenPartial(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 3, 1)
	regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)
	regionResultFor(t, check.ID, "us-east", "ERR", 4*time.Second)
	regionResultFor(t, check.ID, "eu-west", "ERR", 3*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 1, n)

	got := loadCheck(t, check.ID)
	assert.Equal(t, "DEGRADED", got.State,
		"neither OK nor ERR reaches quorum=3 → state=DEGRADED")
	assert.EqualValues(t, 0, countPendingResults(t, check.ID))
}

// TestAggregator_NotEnoughRegionsAlsoDegraded covers the single-region-
// only-delivered case: with RequireQuorum=3 and only 1 result buffered
// (and it aged past the window), the rule still says "neither side
// reached quorum" → DEGRADED. This is the documented behavior for slow
// regions that never report in time.
func TestAggregator_NotEnoughRegionsAlsoDegraded(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 3, 1)
	regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 1, n)

	got := loadCheck(t, check.ID)
	assert.Equal(t, "DEGRADED", got.State,
		"single OK row vs quorum=3 → DEGRADED (below quorum on both sides)")
}

// TestAggregator_NoResultsLeavesStateAlone is the "special case" from
// the spec: when the aggregator tick finds no eligible rows for a
// check, Check.State must NOT change. Pre-set the check to OK and
// verify it stays OK.
func TestAggregator_NoResultsLeavesStateAlone(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 1)
	// Pre-set state and a previous LastEnd so we can detect any
	// accidental overwrite.
	prevEnd := time.Now().Add(-time.Hour)
	require.NoError(t, models.DB().Model(&models.Check{}).
		Where("id = ?", check.ID).
		Updates(map[string]interface{}{
			"state":    "OK",
			"last_end": prevEnd,
			"last_ok":  prevEnd,
		}).Error)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 0, n, "no eligible rows → nothing aggregated")

	got := loadCheck(t, check.ID)
	assert.Equal(t, "OK", got.State, "state must not change with zero eligible rows")
	assert.WithinDuration(t, prevEnd, *got.LastEnd, time.Second,
		"last_end must not be touched when there are no eligible rows")
}

// TestAggregator_MultipleChecksIndependent verifies that a single tick
// processes every check with pending results, not just the first one.
// Two checks, each with 2 regions, each should flip to OK after the
// tick.
func TestAggregator_MultipleChecksIndependent(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, c1 := seedAggregatorWorld(t, 2, 1)
	_, c2 := seedAggregatorWorld(t, 2, 1)

	regionResultFor(t, c1.ID, "ru-msk", "OK", 5*time.Second)
	regionResultFor(t, c1.ID, "us-east", "OK", 4*time.Second)
	regionResultFor(t, c2.ID, "ru-msk", "OK", 5*time.Second)
	regionResultFor(t, c2.ID, "eu-west", "OK", 4*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 2, n, "both checks aggregated in the same tick")

	assert.Equal(t, "OK", loadCheck(t, c1.ID).State)
	assert.Equal(t, "OK", loadCheck(t, c2.ID).State)
	assert.EqualValues(t, 0, countPendingResults(t, c1.ID))
	assert.EqualValues(t, 0, countPendingResults(t, c2.ID))
}

// TestAggregator_AlreadyAggregatedRowsSkipped pins the idempotency
// story: a second tick with no new rows must be a no-op. We pre-mark
// the rows aggregated_at=NOW() and verify the tick returns (0, nil)
// without touching Check.State.
func TestAggregator_AlreadyAggregatedRowsSkipped(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 1)
	r1 := regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)
	r2 := regionResultFor(t, check.ID, "us-east", "OK", 4*time.Second)
	// Pretend a previous aggregator pass already stamped them.
	now := time.Now()
	require.NoError(t, models.DB().Model(&models.CheckRegionResult{}).
		Where("id IN ?", []int64{r1.ID, r2.ID}).
		UpdateColumns(map[string]interface{}{"aggregated_at": now}).Error)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 0, n, "no candidate checks → 0 aggregated")

	got := loadCheck(t, check.ID)
	assert.Equal(t, "UNK", got.State, "already-aggregated rows must not cause a re-decision")
}

// TestAggregator_IgnoresRowsInsideWindow verifies the watermark: rows
// whose CreatedAt is NEWER than (NOW() - window) are NOT eligible and
// must NOT be stamped. With AggregationWindowSeconds=10 and rows aged
// only 2s, the aggregator finds nothing to do.
func TestAggregator_IgnoresRowsInsideWindow(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 10)
	regionResultFor(t, check.ID, "ru-msk", "OK", 2*time.Second)
	regionResultFor(t, check.ID, "us-east", "OK", 1*time.Second)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 0, n, "rows still inside window → no aggregation")

	got := loadCheck(t, check.ID)
	assert.Equal(t, "UNK", got.State)
	assert.EqualValues(t, 2, countPendingResults(t, check.ID),
		"rows inside window stay unaggregated for the next tick")
}

// TestAggregator_SkipsChecksWithQuorumOne guards the candidate SELECT
// filter: even though CheckRegionResult rows are written for
// RequireQuorum=1 checks (via StoreCheckRegionResult), the aggregator
// must not re-decide their state because ApplyRemoteCheckResult
// already did. We simulate by inserting a region row with aggregated_at
// NULL for a quorum=1 check and verifying the tick ignores it.
func TestAggregator_SkipsChecksWithQuorumOne(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 1, 1)
	r := regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)
	require.Nil(t, r.AggregatedAt)

	n, err := models.CheckAggregatorTick()
	require.NoError(t, err)
	assert.Equal(t, 0, n, "quorum=1 checks must be filtered out by the candidate SELECT")

	// The row must stay unaggregated too — the aggregator has no
	// business stamping it.
	assert.EqualValues(t, 1, countPendingResults(t, check.ID))
}

// ---------------------------------------------------------------------------
// StartCheckAggregator: ticker smoke test.
// ---------------------------------------------------------------------------

// TestStartCheckAggregator_TickerFiresOnce is the smoke test for the
// background helper: spin up the aggregator with a tight 10ms ticker
// and a cancellable context, wait for one tick to flip a seeded
// check's state, then cancel so the goroutine exits cleanly. Mirrors
// TestStartDeadWorkerReaper_TickerFiresOnce in shape.
func TestStartCheckAggregator_TickerFiresOnce(t *testing.T) {
	models.Drop()
	models.Migrate()

	_, check := seedAggregatorWorld(t, 2, 1)
	regionResultFor(t, check.ID, "ru-msk", "OK", 5*time.Second)
	regionResultFor(t, check.ID, "us-east", "OK", 4*time.Second)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	models.StartCheckAggregator(ctx, 10*time.Millisecond)

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		got := loadCheck(t, check.ID)
		if got.State == "OK" {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("aggregator goroutine did not flip check to OK within 2s; state=%q", loadCheck(t, check.ID).State)
}

// silence unused import warnings when individual helpers are inlined by
// editors — the package-level references below keep the imports live.
var (
	_ = gorm.ErrRecordNotFound
)