chore: fix cleanup audit log test

internal/jimm/auditlog/auditlog.go

@@ -19,17 +19,25 @@ import (
 	ofganames "github.com/canonical/jimm/v3/internal/openfga/names"
 )
 
+// PollTimeOfDay holds the time hour, minutes and seconds to poll for cleanup.
+type PollTimeOfDay struct {
+	Hours   int
+	Minutes int
+	Seconds int
+}
+
 // auditLogManager provides a means to manage audit logs within JIMM.
 type auditLogManager struct {
 	store                 *db.Database
 	authSvc               *openfga.OFGAClient
 	jimmTag               names.ControllerTag
 	retentionPeriodInDays int
+	pollTimeOfDay         PollTimeOfDay
 }
 
 // NewAuditLogManager returns a new auditLog manager that provides audit Log
 // creation, and removal.
-func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag names.ControllerTag, retentionDays int) (*auditLogManager, error) {
+func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag names.ControllerTag, retentionDays int, pollTime *PollTimeOfDay) (*auditLogManager, error) {
 	if store == nil {
 		return nil, errors.E("auditlog store cannot be nil")
 	}
@@ -39,7 +47,14 @@ func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag
 	if jimmTag.String() == "" {
 		return nil, errors.E("auditlog jimm tag cannot be empty")
 	}
-	return &auditLogManager{store, authSvc, jimmTag, retentionDays}, nil
+	// By default we poll at 9 AM.
+	defaultPollTime := PollTimeOfDay{
+		Hours: 9,
+	}
+	if pollTime == nil {
+		pollTime = &defaultPollTime
+	}
+	return &auditLogManager{store, authSvc, jimmTag, retentionDays, *pollTime}, nil
 }
 
 // addAuditLogEntry causes an entry to be added the the audit log.

internal/jimm/auditlog/auditlog_test.go

@@ -45,7 +45,7 @@ func (s *auditLogManagerSuite) Init(c *qt.C) {
 
 	s.jimmTag = names.NewControllerTag("foo")
 
-	s.manager, err = auditlog.NewAuditLogManager(db, ofgaClient, s.jimmTag, 1)
+	s.manager, err = auditlog.NewAuditLogManager(db, ofgaClient, s.jimmTag, 1, nil)
 	c.Assert(err, qt.IsNil)
 
 	// Create test identity

internal/jimm/auditlog/cleanup.go

@@ -9,55 +9,41 @@ import (
 	"go.uber.org/zap"
 )
 
-// pollTimeOfDay holds the time hour, minutes and seconds to poll at.
-type pollTimeOfDay struct {
-	Hours   int
-	Minutes int
-	Seconds int
-}
-
-var pollDuration = pollTimeOfDay{
-	Hours: 9,
-}
-
-// StartCleanup starts a routine which checks daily for any logs
-// needed to be cleaned up.
+// StartCleanup loop forever and checks daily for any logs
+// that need to be cleaned up. This method should be run
+// in a separate Go routine to avoid blocking, it will terminate
+// when the provided context is cancelled.
 func (j *auditLogManager) StartCleanup(ctx context.Context) {
 	if j.retentionPeriodInDays == 0 {
 		return
 	}
-	go j.poll(ctx)
-}
-
-// poll is designed to be run in a routine where it can be cancelled safely
-// from the service's context. It calculates the poll duration at 9am each day
-// UTC.
-func (j *auditLogManager) poll(ctx context.Context) {
-
 	for {
 		select {
-		case <-time.After(calculateNextPollDuration(time.Now().UTC())):
-			retentionDate := time.Now().AddDate(0, 0, -(j.retentionPeriodInDays))
-			deleted, err := j.store.DeleteAuditLogsBefore(ctx, retentionDate)
-			if err != nil {
-				zapctx.Error(ctx, "failed to cleanup audit logs", zap.Error(err))
-				continue
-			}
-			zapctx.Debug(ctx, "audit log cleanup run successfully", zap.Int64("count", deleted))
+		case <-time.After(calculateNextPollDuration(j.pollTimeOfDay, time.Now().UTC())):
+			j.cleanup(ctx)
 		case <-ctx.Done():
 			zapctx.Debug(ctx, "exiting audit log cleanup polling")
 			return
 		}
 	}
 }
 
+func (j *auditLogManager) cleanup(ctx context.Context) {
+	retentionDate := time.Now().AddDate(0, 0, -(j.retentionPeriodInDays))
+	deleted, err := j.store.DeleteAuditLogsBefore(ctx, retentionDate)
+	if err != nil {
+		zapctx.Error(ctx, "failed to cleanup audit logs", zap.Error(err))
+	}
+	zapctx.Debug(ctx, "audit log cleanup run successfully", zap.Int64("count", deleted))
+}
+
 // calculateNextPollDuration returns the next duration to poll on.
 // We recalculate each time and not rely on running every 24 hours
 // for absolute consistency within ns apart.
-func calculateNextPollDuration(startingTime time.Time) time.Duration {
+func calculateNextPollDuration(pollTime PollTimeOfDay, startingTime time.Time) time.Duration {
 	now := startingTime
-	pollTime := time.Date(now.Year(), now.Month(), now.Day(), pollDuration.Hours, pollDuration.Minutes, pollDuration.Seconds, 0, time.UTC)
-	tillNextPoll := pollTime.Sub(now)
+	pollTimeToday := time.Date(now.Year(), now.Month(), now.Day(), pollTime.Hours, pollTime.Minutes, pollTime.Seconds, 0, time.UTC)
+	tillNextPoll := pollTimeToday.Sub(now)
 	var d time.Duration
 	// If the next poll time is behind the current time
 	if tillNextPoll < 0 {

internal/jimm/auditlog/cleanup_test.go

@@ -1,4 +1,5 @@
 // Copyright 2025 Canonical.
+
 package auditlog_test
 
 import (
@@ -7,66 +8,79 @@ import (
 	"time"
 
 	qt "github.com/frankban/quicktest"
+	"github.com/juju/names/v5"
 
+	"github.com/canonical/jimm/v3/internal/db"
 	"github.com/canonical/jimm/v3/internal/dbmodel"
-	"github.com/canonical/jimm/v3/internal/errors"
 	"github.com/canonical/jimm/v3/internal/jimm/auditlog"
+	"github.com/canonical/jimm/v3/internal/testutils/jimmtest"
 )
 
-func (s *auditLogManagerSuite) TestAuditLogCleanupServicePurgesLogs(c *qt.C) {
+func TestAuditLogCleanupServicePurgesLogs(t *testing.T) {
+	c := qt.New(t)
 	c.Parallel()
+
 	ctx := context.Background()
-	now := time.Now().UTC().Round(time.Millisecond)
 
-	err := s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
-		Time: now.AddDate(0, 0, -1),
-	})
-	c.Check(errors.ErrorCode(err), qt.Equals, errors.CodeUpgradeInProgress)
+	db := &db.Database{
+		DB: jimmtest.PostgresDB(c, time.Now),
+	}
+	err := db.Migrate(context.Background())
+	c.Assert(err, qt.IsNil)
+
+	ofgaClient, _, _, err := jimmtest.SetupTestOFGAClient(c.Name())
+	c.Assert(err, qt.IsNil)
+
+	jimmTag := names.NewControllerTag("foo")
+
+	manager, err := auditlog.NewAuditLogManager(db, ofgaClient, jimmTag, 1, nil)
+	c.Assert(err, qt.IsNil)
+
+	now := time.Now().UTC()
+
+	// A log from today
+	c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
+		Time: now.AddDate(0, 0, 0),
+	}), qt.IsNil)
 
 	// A log from 1 day ago
-	c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
+	c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
 		Time: now.AddDate(0, 0, -1),
 	}), qt.IsNil)
 
 	// A log from 2 days ago
-	c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
+	c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
 		Time: now.AddDate(0, 0, -2),
 	}), qt.IsNil)
 
-	// A log from 3 days ago
-	c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{
-		Time: now.AddDate(0, 0, -3),
-	}), qt.IsNil)
-
 	// Check 3 created
 	logs := make([]dbmodel.AuditLogEntry, 0)
-	err = s.db.DB.Find(&logs).Error
+	err = db.DB.Find(&logs).Error
 	c.Assert(err, qt.IsNil)
 	c.Assert(logs, qt.HasLen, 3)
 
-	auditlog.PollDuration.Hours = now.Hour()
-	auditlog.PollDuration.Minutes = now.Minute()
-	auditlog.PollDuration.Seconds = now.Second() + 2
-	// Suite is setup to clean logs more than 1 day old.
-	s.manager.StartCleanup(ctx)
+	// Manager is setup above to remove logs older than 1 day.
+	manager.Cleanup(ctx)
 
 	// Check 2 were purged
 	logs = make([]dbmodel.AuditLogEntry, 0)
-	err = s.db.DB.Find(&logs).Error
+	err = db.DB.Find(&logs).Error
 	c.Assert(err, qt.IsNil)
-	c.Assert(logs, qt.HasLen, 3)
+	c.Assert(logs, qt.HasLen, 1)
 }
 
 func TestCalculateNextPollDuration(t *testing.T) {
 	c := qt.New(t)
 
+	pollTime := auditlog.PollTimeOfDay{Hours: 9}
+
 	// Test where 9am is behind 12pm
 	startingTime := time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC)
-	d := auditlog.CalculateNextPollDuration(startingTime)
+	d := auditlog.CalculateNextPollDuration(pollTime, startingTime)
 	c.Assert(d, qt.Equals, time.Hour*21)
 
-	// Test where 9am is ahead of 7pm
+	// Test where 9am is ahead of 7am
 	startingTime = time.Date(2023, 1, 1, 7, 0, 0, 0, time.UTC)
-	d = auditlog.CalculateNextPollDuration(startingTime)
+	d = auditlog.CalculateNextPollDuration(pollTime, startingTime)
 	c.Assert(d, qt.Equals, time.Hour*2)
 }

internal/jimm/auditlog/export_test.go

@@ -2,10 +2,15 @@
 
 package auditlog
 
+import "context"
+
 // AuditLogManager is a type alias to export auditLogManager for use in tests.
 type AuditLogManager = auditLogManager
 
 var (
-	PollDuration              = pollDuration
 	CalculateNextPollDuration = calculateNextPollDuration
 )
+
+func (j *auditLogManager) Cleanup(ctx context.Context) {
+	j.cleanup(ctx)
+}

internal/jimm/jimm.go

@@ -300,7 +300,7 @@ func New(p Parameters) (*JIMM, error) {
 	}
 	j.loginManager = loginManager
 
-	auditLogManager, err := auditlog.NewAuditLogManager(j.Database, j.OpenFGAClient, j.ResourceTag(), p.AuditLogRetentionDays)
+	auditLogManager, err := auditlog.NewAuditLogManager(j.Database, j.OpenFGAClient, j.ResourceTag(), p.AuditLogRetentionDays, nil)
 	if err != nil {
 		return nil, err
 	}