test: fix flaky audit log and service account status tests

The service account status test indexed PublicKeys[0] right after a non-fatal assert.Len. A slow reconcile that left the slice empty would panic instead of retrying, so the index access now sits behind the length check and the poller waits for the key to be aggregated.

The two audit log tests seeded 1500 rows as separate autocommit transactions. Under full-suite disk contention those commits triggered enough checkpoint fsyncs to overrun the 30s context deadline mid-statement, so seeding now runs inside a single transaction.

Signed-off-by: Oguz Kilcan <oguz.kilcan@siderolabs.com>

Author: oguzkilcan

internal/backend/runtime/omni/audit/auditlog/auditlogsqlite/auditlogsqlite_test.go

@@ -313,23 +313,56 @@ func TestRemoveByMaxSizeBatchCap(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 
 	// Seed 1500 events into a store with no size limit.
-	seedStore, db := setupStore(ctx, t, logger)
+	_, db := setupStore(ctx, t, logger)
 
 	const totalEvents = 1500
 
-	for i := range totalEvents {
-		evt := auditlog.Event{
-			Type:       "batch-cap-test",
-			TimeMillis: int64(i) * 1000,
-			Data: &auditlog.Data{
+	require.NoError(t, func() (err error) {
+		var conn *zombiesqlite.Conn
+
+		if conn, err = db.Take(ctx); err != nil {
+			return err
+		}
+
+		defer db.Put(conn)
+
+		endTx, err := sqlitex.ImmediateTransaction(conn)
+		if err != nil {
+			return err
+		}
+
+		defer endTx(&err)
+
+		insertQuery := fmt.Sprintf("INSERT INTO %s (%s, %s, %s) VALUES ($type, $ts, $data)",
+			auditlogsqlite.TableName, "event_type", "event_ts_ms", "event_data")
+
+		for i := range totalEvents {
+			data, mErr := json.Marshal(&auditlog.Data{
 				Session: auditlog.Session{
 					UserID: fmt.Sprintf("user-%d", i),
 					Email:  fmt.Sprintf("user-%d@example.com", i),
 				},
-			},
+			})
+			if mErr != nil {
+				return mErr
+			}
+
+			q, qErr := sqlitexx.NewQuery(conn, insertQuery)
+			if qErr != nil {
+				return qErr
+			}
+
+			if err = q.
+				BindString("$type", "batch-cap-test").
+				BindInt64("$ts", int64(i)*1000).
+				BindBytes("$data", data).
+				Exec(); err != nil {
+				return err
+			}
 		}
-		require.NoError(t, seedStore.Write(ctx, evt))
-	}
+
+		return nil
+	}())
 
 	// Create a new store on the same DB with maxSize=1 (effectively 0) and probability=1.0.
 	// This means all rows exceed the limit, but the batch cap (1000) should prevent
@@ -386,30 +419,48 @@ func TestRemoveBatching(t *testing.T) {
 	rangeStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
 	rangeEnd := time.Date(2026, 1, 2, 0, 0, 0, 0, time.UTC)
 
-	conn, err := db.Take(ctx)
-	require.NoError(t, err)
+	require.NoError(t, func() (err error) {
+		var conn *zombiesqlite.Conn
 
-	t.Cleanup(func() {
-		db.Put(conn)
-	})
+		if conn, err = db.Take(ctx); err != nil {
+			return err
+		}
 
-	for i := range inRangeCount {
-		ts := rangeStart.Add(time.Duration(i) * time.Second).UnixMilli()
+		defer db.Put(conn)
 
-		q, qErr := sqlitexx.NewQuery(conn,
-			fmt.Sprintf("INSERT INTO %s (%s, %s) VALUES ($ts, $data)", auditlogsqlite.TableName, "event_ts_ms", "event_data"))
-		require.NoError(t, qErr)
-		require.NoError(t, q.BindInt64("$ts", ts).BindBytes("$data", []byte(`{}`)).Exec())
-	}
+		endTx, err := sqlitex.ImmediateTransaction(conn)
+		if err != nil {
+			return err
+		}
 
-	for i := range outRangeCount {
-		ts := rangeEnd.Add(time.Duration(i+1) * time.Hour).UnixMilli()
+		defer endTx(&err)
 
-		q, qErr := sqlitexx.NewQuery(conn,
-			fmt.Sprintf("INSERT INTO %s (%s, %s) VALUES ($ts, $data)", auditlogsqlite.TableName, "event_ts_ms", "event_data"))
-		require.NoError(t, qErr)
-		require.NoError(t, q.BindInt64("$ts", ts).BindBytes("$data", []byte(`{}`)).Exec())
-	}
+		insertQuery := fmt.Sprintf("INSERT INTO %s (%s, %s) VALUES ($ts, $data)",
+			auditlogsqlite.TableName, "event_ts_ms", "event_data")
+
+		insert := func(ts int64) error {
+			q, qErr := sqlitexx.NewQuery(conn, insertQuery)
+			if qErr != nil {
+				return qErr
+			}
+
+			return q.BindInt64("$ts", ts).BindBytes("$data", []byte(`{}`)).Exec()
+		}
+
+		for i := range inRangeCount {
+			if err = insert(rangeStart.Add(time.Duration(i) * time.Second).UnixMilli()); err != nil {
+				return err
+			}
+		}
+
+		for i := range outRangeCount {
+			if err = insert(rangeEnd.Add(time.Duration(i+1) * time.Hour).UnixMilli()); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}())
 
 	// Remove all events in range — requires multiple batches (1000 + 500).
 	require.NoError(t, store.Remove(ctx, rangeStart, rangeEnd))

internal/backend/runtime/omni/controllers/omni/service_account_status_test.go

@@ -69,7 +69,10 @@ func (suite *ClusterServiceAccountStatusSuite) TestReconcile() {
 
 	rtestutils.AssertResources(ctx, suite.T(), suite.state, []string{serviceAccount.Metadata().ID()}, func(res *auth.ServiceAccountStatus, assert *assert.Assertions) {
 		assert.Equal(string(role.Admin), res.TypedSpec().Value.Role)
-		assert.Len(res.TypedSpec().Value.PublicKeys, 1)
+
+		if !assert.Len(res.TypedSpec().Value.PublicKeys, 1) {
+			return
+		}
 
 		if assert.NotNil(res.TypedSpec().Value.PublicKeys[0].Created) {
 			assert.False(res.TypedSpec().Value.PublicKeys[0].Created.AsTime().IsZero())
@@ -86,7 +89,9 @@ func (suite *ClusterServiceAccountStatusSuite) TestReconcile() {
 	require.NoError(suite.state.Create(suite.ctx, pkLastActive))
 
 	rtestutils.AssertResources(ctx, suite.T(), suite.state, []string{serviceAccount.Metadata().ID()}, func(res *auth.ServiceAccountStatus, assert *assert.Assertions) {
-		assert.Len(res.TypedSpec().Value.PublicKeys, 1)
+		if !assert.Len(res.TypedSpec().Value.PublicKeys, 1) {
+			return
+		}
 
 		if assert.NotNil(res.TypedSpec().Value.PublicKeys[0].LastUsed) {
 			assert.WithinDuration(time.Now(), res.TypedSpec().Value.PublicKeys[0].LastUsed.AsTime(), 5*time.Second)