Add usage statistics to measure policy violations (#8199)

server/datastore/mysql/policies.go

@@ -3,6 +3,7 @@ package mysql
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"sort"
 	"strings"
@@ -690,3 +691,156 @@ func (ds *Datastore) CleanupPolicyMembership(ctx context.Context, now time.Time)
 
 	return nil
 }
+
+// PolicyViolationDays is a structure used for aggregate counts of policy violation days.
+type PolicyViolationDays struct {
+	// FailingHostCount is an aggregate count of actual policy violations days. One actual policy
+	// violation day is added for each policy that a host is failing at the time of the count.
+	FailingHostCount uint `json:"failing_host_count" db:"failing_host_count"`
+	// TotalHostCount is an aggregate count of possible policy violations days. One possible policy
+	// violation day is added for each policy that a host is a member of at the time of the count.
+	TotalHostCount uint `json:"total_host_count" db:"total_host_count"`
+}
+
+func (ds *Datastore) IncrementPolicyViolationDays(ctx context.Context) error {
+	return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
+		return incrementViolationDaysDB(ctx, tx)
+	})
+}
+
+func incrementViolationDaysDB(ctx context.Context, tx sqlx.ExtContext) error {
+	const (
+		statsID        = 0
+		statsType      = "policy_violation_days"
+		updateInterval = 24 * time.Hour
+	)
+
+	var prevFailing uint
+	var prevTotal uint
+	var shouldIncrement bool
+
+	// get current count of policy violation days from `aggregated_stats``
+	selectStmt := `
+		SELECT 
+			json_value, 
+			created_at, 
+			updated_at 
+		FROM 
+			aggregated_stats 
+		WHERE 
+			id = ? AND type = ?`
+	dest := struct {
+		CreatedAt time.Time       `json:"created_at" db:"created_at"`
+		UpdatedAt time.Time       `json:"updated_at" db:"updated_at"`
+		StatsJSON json.RawMessage `json:"json_value" db:"json_value"`
+	}{}
+
+	err := sqlx.GetContext(ctx, tx, &dest, selectStmt, statsID, statsType)
+	switch {
+	case err == sql.ErrNoRows:
+		// no previous counts exists so initialize counts as zero and proceed to increment
+		prevFailing = 0
+		prevTotal = 0
+		shouldIncrement = true
+	case err != nil:
+		return ctxerr.Wrap(ctx, err, "selecting policy violation days aggregated stats")
+	default:
+		// increment previous counts if interval has elapsed
+		var prevStats PolicyViolationDays
+		if err := json.Unmarshal(dest.StatsJSON, &prevStats); err != nil {
+			return ctxerr.Wrap(ctx, err, "unmarshal policy violation counts")
+		}
+		prevFailing = prevStats.FailingHostCount
+		prevTotal = prevStats.TotalHostCount
+		shouldIncrement = time.Now().After(dest.UpdatedAt.Add(updateInterval))
+	}
+
+	if !shouldIncrement {
+		return nil
+	}
+
+	// increment count of policy violation days by total number of failing records from
+	// `policy_membership`
+	var newCounts PolicyViolationDays
+	if err := sqlx.GetContext(ctx, tx, &newCounts, `
+	     SELECT	(select count(*) from policy_membership where passes=0) as failing_host_count,
+       		(select count(*) from policy_membership) as total_host_count`,
+	); err != nil {
+		return ctxerr.Wrap(ctx, err, "count policy violation days")
+	}
+	newCounts.FailingHostCount = prevFailing + newCounts.FailingHostCount
+	newCounts.TotalHostCount = prevTotal + newCounts.TotalHostCount
+	statsJSON, err := json.Marshal(newCounts)
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "marshal policy violation counts")
+	}
+
+	// upsert `aggregated_stats` with new count
+	upsertStmt := `
+		INSERT INTO 
+			aggregated_stats (id, type, json_value) 
+		VALUES (?, ?, ?)
+		ON DUPLICATE KEY UPDATE 
+			json_value = VALUES(json_value)`
+	if _, err := tx.ExecContext(ctx, upsertStmt, statsID, statsType, statsJSON); err != nil {
+		return ctxerr.Wrap(ctx, err, "update policy violation days aggregated stats")
+	}
+
+	return nil
+}
+
+func (ds *Datastore) InitializePolicyViolationDays(ctx context.Context) error {
+	return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
+		return initializePolicyViolationDaysDB(ctx, tx)
+	})
+}
+
+func initializePolicyViolationDaysDB(ctx context.Context, tx sqlx.ExtContext) error {
+	const (
+		statsID   = 0
+		statsType = "policy_violation_days"
+	)
+
+	statsJSON, err := json.Marshal(PolicyViolationDays{})
+	if err != nil {
+		return ctxerr.Wrap(ctx, err, "marshal policy violation counts")
+	}
+
+	stmt := `
+		INSERT INTO 
+			aggregated_stats (id, type, json_value) 
+		VALUES (?, ?, ?)
+		ON DUPLICATE KEY UPDATE 
+			json_value = VALUES(json_value),
+			created_at = CURRENT_TIMESTAMP`
+	if _, err := tx.ExecContext(ctx, stmt, statsID, statsType, statsJSON); err != nil {
+		return ctxerr.Wrap(ctx, err, "initialize policy violation days aggregated stats")
+	}
+
+	return nil
+}
+
+func amountPolicyViolationDaysDB(ctx context.Context, tx sqlx.QueryerContext) (int, int, error) {
+	const (
+		statsID   = 0
+		statsType = "policy_violation_days"
+	)
+	var statsJSON json.RawMessage
+	if err := sqlx.GetContext(ctx, tx, &statsJSON, `
+		SELECT 
+			json_value 
+		FROM 
+			aggregated_stats 
+		WHERE 
+			id = ? AND type = ?
+	`, statsID, statsType); err != nil {
+		return 0, 0, err
+	}
+
+	var counts PolicyViolationDays
+	if err := json.Unmarshal(statsJSON, &counts); err != nil {
+		return 0, 0, ctxerr.Wrap(ctx, err, "unmarshal policy violation counts")
+	}
+
+	return int(counts.FailingHostCount), int(counts.TotalHostCount), nil
+}

server/datastore/mysql/policies_test.go

@@ -44,6 +44,7 @@ func TestPolicies(t *testing.T) {
 		{"PlatformUpdate", testPolicyPlatformUpdate},
 		{"CleanupPolicyMembership", testPolicyCleanupPolicyMembership},
 		{"DeleteAllPolicyMemberships", testDeleteAllPolicyMemberships},
+		{"PolicyViolationDays", testPolicyViolationDays},
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
@@ -1771,6 +1772,105 @@ func assertPolicyMembership(t *testing.T, ds *Datastore, polsByName map[string]*
 	}
 }
 
+func testPolicyViolationDays(t *testing.T, ds *Datastore) {
+	ctx := context.Background()
+	then := time.Now().Add(-48 * time.Hour)
+
+	setStatsTimestampDB := func(updatedAt time.Time) error {
+		_, err := ds.writer.ExecContext(ctx, `
+			UPDATE aggregated_stats SET created_at = ?, updated_at = ? WHERE id = ? AND type = ?
+		`, then, updatedAt, 0, "policy_violation_days")
+		return err
+	}
+
+	user := test.NewUser(t, ds, "Bob", "bob@example.com", true)
+
+	hosts := make([]*fleet.Host, 3)
+	for i, name := range []string{"h1", "h2", "h3"} {
+		id := fmt.Sprintf("%s-%d", strings.ReplaceAll(t.Name(), "/", "_"), i)
+		h, err := ds.NewHost(ctx, &fleet.Host{
+			OsqueryHostID:   id,
+			DetailUpdatedAt: then,
+			LabelUpdatedAt:  then,
+			PolicyUpdatedAt: then,
+			SeenTime:        then,
+			NodeKey:         id,
+			UUID:            id,
+			Hostname:        name,
+		})
+		require.NoError(t, err)
+		hosts[i] = h
+	}
+
+	createPolStmt := `INSERT INTO policies (name, query, description, author_id, platforms, created_at, updated_at) VALUES (?, ?, '', ?, ?, ?, ?)`
+	res, err := ds.writer.ExecContext(ctx, createPolStmt, "test_pol", "select 1", user.ID, "", then, then)
+	require.NoError(t, err)
+	id, _ := res.LastInsertId()
+	pol, err := ds.Policy(ctx, uint(id))
+	require.NoError(t, err)
+
+	require.NoError(t, ds.InitializePolicyViolationDays(ctx)) // sets starting violation count to zero
+
+	// initialize policy statuses: 1 failling, 2 passing
+	require.NoError(t, ds.RecordPolicyQueryExecutions(context.Background(), hosts[0], map[uint]*bool{pol.ID: ptr.Bool(false)}, then, false))
+	require.NoError(t, ds.RecordPolicyQueryExecutions(context.Background(), hosts[1], map[uint]*bool{pol.ID: ptr.Bool(true)}, then, false))
+	require.NoError(t, ds.RecordPolicyQueryExecutions(context.Background(), hosts[2], map[uint]*bool{pol.ID: ptr.Bool(true)}, then, false))
+
+	// setup db for test: starting counts zero, more than 24h since last updated, one outstanding violation
+	require.NoError(t, setStatsTimestampDB(time.Now().Add(-25*time.Hour)))
+	require.NoError(t, ds.IncrementPolicyViolationDays(ctx))
+	actual, possible, err := amountPolicyViolationDaysDB(ctx, ds.reader)
+	require.NoError(t, err)
+	// actual should increment from 0 -> 1 (+1 outstanding violation)
+	require.Equal(t, 1, actual)
+	// possible should increment from 0 -> 3 (3 total hosts * 1 policy)
+	require.Equal(t, 3, possible)
+	// reset violation counts to zero for next test
+	require.NoError(t, ds.InitializePolicyViolationDays(ctx))
+
+	// setup for test: starting counts zero, less than 24h since last updated, one outstanding violation
+	require.NoError(t, setStatsTimestampDB(time.Now().Add(-1*time.Hour)))
+	require.NoError(t, ds.IncrementPolicyViolationDays(ctx))
+	actual, possible, err = amountPolicyViolationDaysDB(ctx, ds.reader)
+	require.NoError(t, err)
+	// count should not increment from zero
+	require.Equal(t, 0, actual)
+	// possible should not increment from zero
+	require.Equal(t, 0, possible)
+	// leave counts at zero for next test
+
+	// setup for test: starting count zero, more than 24h since last updated, add second outstanding violation
+	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, hosts[1], map[uint]*bool{pol.ID: ptr.Bool(false)}, time.Now(), false))
+	require.NoError(t, setStatsTimestampDB(time.Now().Add(-25*time.Hour)))
+	require.NoError(t, ds.IncrementPolicyViolationDays(ctx))
+	actual, possible, err = amountPolicyViolationDaysDB(ctx, ds.reader)
+	require.NoError(t, err)
+	// actual should increment from 0 -> 2 (+2 outstanding violations)
+	require.Equal(t, 2, actual) // leave count at two for next test
+	// possible should increment from 0 -> 3 (3 total hosts * 1 policy)
+	require.Equal(t, 3, possible)
+	// leave counts at 2 actual and 3 possible for next test
+
+	// setup for test: starting counts at 2 actual and 3 possible, more than 24h since last updated, resolve one outstaning violation
+	require.NoError(t, ds.RecordPolicyQueryExecutions(ctx, hosts[1], map[uint]*bool{pol.ID: ptr.Bool(true)}, time.Now(), false))
+	require.NoError(t, setStatsTimestampDB(time.Now().Add(-25*time.Hour)))
+	require.NoError(t, ds.IncrementPolicyViolationDays(ctx))
+	actual, possible, err = amountPolicyViolationDaysDB(ctx, ds.reader)
+	require.NoError(t, err)
+	// actual should increment from 2 -> 3 (+1 outstanding violation)
+	require.Equal(t, 3, actual)
+	// possible should increment from 3 -> 6 (3 total hosts * 1 policy)
+	require.Equal(t, 6, possible)
+	// leave counts at 3 actual and 6 possible
+
+	// attempt again immediately after last update, counts should not increment
+	require.NoError(t, ds.IncrementPolicyViolationDays(ctx))
+	actual, possible, err = amountPolicyViolationDaysDB(ctx, ds.reader)
+	require.NoError(t, err)
+	require.Equal(t, 3, actual)
+	require.Equal(t, 6, possible)
+}
+
 func testPolicyCleanupPolicyMembership(t *testing.T, ds *Datastore) {
 	ctx := context.Background()
 	user := test.NewUser(t, ds, "Bob", "bob@example.com", true)

server/datastore/mysql/statistics.go

@@ -9,6 +9,7 @@ import (
 	"github.com/fleetdm/fleet/v4/server/config"
 	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
 	"github.com/fleetdm/fleet/v4/server/fleet"
+	"github.com/go-kit/kit/log/level"
 	"github.com/jmoiron/sqlx"
 	"github.com/kolide/kit/version"
 )
@@ -48,6 +49,12 @@ func (ds *Datastore) ShouldSendStatistics(ctx context.Context, frequency time.Du
 		if err != nil {
 			return ctxerr.Wrap(ctx, err, "amount active users")
 		}
+		amountPolicyViolationDaysActual, amountPolicyViolationDaysPossible, err := amountPolicyViolationDaysDB(ctx, ds.writer)
+		if err == sql.ErrNoRows {
+			level.Debug(ds.logger).Log("msg", "amount policy violation days", "err", err)
+		} else if err != nil {
+			return ctxerr.Wrap(ctx, err, "amount policy violation days")
+		}
 		storedErrs, err := ctxerr.Aggregate(ctx)
 		if err != nil {
 			return ctxerr.Wrap(ctx, err, "statistics error store")
@@ -67,6 +74,8 @@ func (ds *Datastore) ShouldSendStatistics(ctx context.Context, frequency time.Du
 		stats.SystemUsersEnabled = appConfig.Features.EnableHostUsers
 		stats.HostsStatusWebHookEnabled = appConfig.WebhookSettings.HostStatusWebhook.Enable
 		stats.NumWeeklyActiveUsers = amountWeeklyUsers
+		stats.NumWeeklyPolicyViolationDaysActual = amountPolicyViolationDaysActual
+		stats.NumWeeklyPolicyViolationDaysPossible = amountPolicyViolationDaysPossible
 		stats.HostsEnrolledByOperatingSystem = enrolledHostsByOS
 		stats.StoredErrors = storedErrs
 		stats.NumHostsNotResponding = amountHostsNotResponding
@@ -129,3 +138,11 @@ func (ds *Datastore) RecordStatisticsSent(ctx context.Context) error {
 	_, err := ds.writer.ExecContext(ctx, `UPDATE statistics SET updated_at = CURRENT_TIMESTAMP LIMIT 1`)
 	return ctxerr.Wrap(ctx, err, "update statistics")
 }
+
+func (ds *Datastore) CleanupStatistics(ctx context.Context) error {
+	// reset weekly count of policy violation days
+	if err := ds.InitializePolicyViolationDays(ctx); err != nil {
+		return err
+	}
+	return nil
+}

server/datastore/mysql/statistics_test.go

@@ -113,6 +113,20 @@ func testStatisticsShouldSend(t *testing.T, ds *Datastore) {
 			OrgLogoURL: "localhost:8080/logo.png",
 		},
 	})
+	require.NoError(t, err)
+
+	// Initialize policy violation days for test
+	pvdJSON, err := json.Marshal(PolicyViolationDays{FailingHostCount: 5, TotalHostCount: 10})
+	require.NoError(t, err)
+	_, err = ds.writer.ExecContext(ctx, `
+		INSERT INTO
+			aggregated_stats (id, type, json_value, created_at, updated_at)
+		VALUES (?, ?, CAST(? AS JSON), ?, ?)
+		ON DUPLICATE KEY UPDATE
+			json_value = VALUES(json_value),
+			updated_at = VALUES(updated_at)
+	`, 0, "policy_violation_days", pvdJSON, time.Now().Add(-48*time.Hour), time.Now().Add(-7*24*time.Hour))
+	require.NoError(t, err)
 
 	require.NoError(t, err)
 	config.Features.EnableSoftwareInventory = false
@@ -146,6 +160,8 @@ func testStatisticsShouldSend(t *testing.T, ds *Datastore) {
 	assert.Equal(t, stats.VulnDetectionEnabled, false)
 	assert.Equal(t, stats.HostsStatusWebHookEnabled, true)
 	assert.Equal(t, stats.NumWeeklyActiveUsers, 1)
+	assert.Equal(t, stats.NumWeeklyPolicyViolationDaysActual, 5)
+	assert.Equal(t, stats.NumWeeklyPolicyViolationDaysPossible, 10)
 	assert.Equal(t, string(stats.StoredErrors), `[{"count":10,"loc":["a","b","c"]}]`)
 
 	firstIdentifier := stats.AnonymousIdentifier
@@ -236,6 +252,7 @@ func testStatisticsShouldSend(t *testing.T, ds *Datastore) {
 	assert.Equal(t, stats.NumUsers, 2)
 	assert.Equal(t, stats.NumWeeklyActiveUsers, 0)          // no active user since last stats were sent
 	require.Len(t, stats.HostsEnrolledByOperatingSystem, 3) // empty platform, rhel and macos
+	assert.Equal(t, stats.NumWeeklyPolicyViolationDaysActual, 5)
 	require.ElementsMatch(t, []fleet.HostsCountByOSVersion{
 		{Version: "Fedora 35", NumEnrolled: 2},
 		{Version: "Fedora 36", NumEnrolled: 1},
@@ -256,6 +273,10 @@ func testStatisticsShouldSend(t *testing.T, ds *Datastore) {
 	_, err = ds.NewSession(ctx, u1.ID, "session_key4")
 	require.NoError(t, err)
 
+	// CleanupStatistics resets policy violation days
+	err = ds.CleanupStatistics(ctx)
+	require.NoError(t, err)
+
 	// wait a bit and resend statistics
 	time.Sleep(1100 * time.Millisecond) // ensure the DB timestamp is not in the same second
 
@@ -268,6 +289,8 @@ func testStatisticsShouldSend(t *testing.T, ds *Datastore) {
 	assert.Equal(t, stats.NumHostsEnrolled, 5)
 	assert.Equal(t, stats.NumUsers, 2)
 	assert.Equal(t, stats.NumWeeklyActiveUsers, 1)
+	assert.Equal(t, stats.NumWeeklyPolicyViolationDaysActual, 0)
+	assert.Equal(t, stats.NumWeeklyPolicyViolationDaysPossible, 0)
 	assert.Equal(t, string(stats.StoredErrors), `[{"count":10,"loc":["a","b","c"]}]`)
 
 	// Add host to test hosts not responding stats

server/fleet/datastore.go

@@ -437,6 +437,9 @@ type Datastore interface {
 
 	ShouldSendStatistics(ctx context.Context, frequency time.Duration, config config.FleetConfig, license *LicenseInfo) (StatisticsPayload, bool, error)
 	RecordStatisticsSent(ctx context.Context) error
+	// CleanupStatistics executes cleanup tasks to be performed upon successful transmission of
+	// statistics.
+	CleanupStatistics(ctx context.Context) error
 
 	///////////////////////////////////////////////////////////////////////////////
 	// GlobalPoliciesStore
@@ -483,6 +486,14 @@ type Datastore interface {
 	TeamPolicy(ctx context.Context, teamID uint, policyID uint) (*Policy, error)
 
 	CleanupPolicyMembership(ctx context.Context, now time.Time) error
+	// IncrementPolicyViolationDays increments the aggregate count of policy violation days. One
+	// policy violation day is added for each policy that a host is failing as of the time the count
+	// is incremented. The count only increments once per 24-hour interval. If the interval has not
+	// elapsed, IncrementPolicyViolationDays returns nil without incrementing the count.
+	IncrementPolicyViolationDays(ctx context.Context) error
+	// InitializePolicyViolationDays sets the aggregated count of policy violation days to zero. If
+	// a record of the count already exists, its `created_at` timestamp is updated to the current timestamp.
+	InitializePolicyViolationDays(ctx context.Context) error
 
 	///////////////////////////////////////////////////////////////////////////////
 	// Locking

server/fleet/statistics.go

@@ -6,22 +6,31 @@ import (
 )
 
 type StatisticsPayload struct {
-	AnonymousIdentifier            string                             `json:"anonymousIdentifier"`
-	FleetVersion                   string                             `json:"fleetVersion"`
-	LicenseTier                    string                             `json:"licenseTier"`
-	Organization                   string                             `json:"organization"`
-	NumHostsEnrolled               int                                `json:"numHostsEnrolled"`
-	NumUsers                       int                                `json:"numUsers"`
-	NumTeams                       int                                `json:"numTeams"`
-	NumPolicies                    int                                `json:"numPolicies"`
-	NumLabels                      int                                `json:"numLabels"`
-	SoftwareInventoryEnabled       bool                               `json:"softwareInventoryEnabled"`
-	VulnDetectionEnabled           bool                               `json:"vulnDetectionEnabled"`
-	SystemUsersEnabled             bool                               `json:"systemUsersEnabled"`
-	HostsStatusWebHookEnabled      bool                               `json:"hostsStatusWebHookEnabled"`
-	NumWeeklyActiveUsers           int                                `json:"numWeeklyActiveUsers"`
-	HostsEnrolledByOperatingSystem map[string][]HostsCountByOSVersion `json:"hostsEnrolledByOperatingSystem"`
-	StoredErrors                   json.RawMessage                    `json:"storedErrors"`
+	AnonymousIdentifier       string `json:"anonymousIdentifier"`
+	FleetVersion              string `json:"fleetVersion"`
+	LicenseTier               string `json:"licenseTier"`
+	Organization              string `json:"organization"`
+	NumHostsEnrolled          int    `json:"numHostsEnrolled"`
+	NumUsers                  int    `json:"numUsers"`
+	NumTeams                  int    `json:"numTeams"`
+	NumPolicies               int    `json:"numPolicies"`
+	NumLabels                 int    `json:"numLabels"`
+	SoftwareInventoryEnabled  bool   `json:"softwareInventoryEnabled"`
+	VulnDetectionEnabled      bool   `json:"vulnDetectionEnabled"`
+	SystemUsersEnabled        bool   `json:"systemUsersEnabled"`
+	HostsStatusWebHookEnabled bool   `json:"hostsStatusWebHookEnabled"`
+	NumWeeklyActiveUsers      int    `json:"numWeeklyActiveUsers"`
+	// NumWeeklyPolicyViolationDaysActual is an aggregate count of actual policy violation days. One
+	// policy violation day is added for each policy that a host is failing as of the time the count
+	// is incremented. The count increments once per 24-hour interval and resets each week.
+	NumWeeklyPolicyViolationDaysActual int `json:"numWeeklyPolicyViolationDaysActual"`
+	// NumWeeklyPolicyViolationDaysActual is an aggregate count of possible policy violation
+	// days. The count is incremented by the organization's total number of policies
+	// mulitplied by the total number of hosts as of the time the count is incremented. The count
+	// increments once per 24-hour interval and resets each week.
+	NumWeeklyPolicyViolationDaysPossible int                                `json:"numWeeklyPolicyViolationDaysPossible"`
+	HostsEnrolledByOperatingSystem       map[string][]HostsCountByOSVersion `json:"hostsEnrolledByOperatingSystem"`
+	StoredErrors                         json.RawMessage                    `json:"storedErrors"`
 	// NumHostsNotResponding is a count of hosts that connect to Fleet successfully but fail to submit results for distributed queries.
 	NumHostsNotResponding int `json:"numHostsNotResponding"`
 }

server/mock/datastore_mock.go

@@ -339,6 +339,8 @@ type ShouldSendStatisticsFunc func(ctx context.Context, frequency time.Duration,
 
 type RecordStatisticsSentFunc func(ctx context.Context) error
 
+type CleanupStatisticsFunc func(ctx context.Context) error
+
 type ApplyPolicySpecsFunc func(ctx context.Context, authorID uint, specs []*fleet.PolicySpec) error
 
 type NewGlobalPolicyFunc func(ctx context.Context, authorID *uint, args fleet.PolicyPayload) (*fleet.Policy, error)
@@ -381,6 +383,10 @@ type TeamPolicyFunc func(ctx context.Context, teamID uint, policyID uint) (*flee
 
 type CleanupPolicyMembershipFunc func(ctx context.Context, now time.Time) error
 
+type IncrementPolicyViolationDaysFunc func(ctx context.Context) error
+
+type InitializePolicyViolationDaysFunc func(ctx context.Context) error
+
 type LockFunc func(ctx context.Context, name string, owner string, expiration time.Duration) (bool, error)
 
 type UnlockFunc func(ctx context.Context, name string, owner string) error
@@ -967,6 +973,9 @@ type DataStore struct {
 	RecordStatisticsSentFunc        RecordStatisticsSentFunc
 	RecordStatisticsSentFuncInvoked bool
 
+	CleanupStatisticsFunc        CleanupStatisticsFunc
+	CleanupStatisticsFuncInvoked bool
+
 	ApplyPolicySpecsFunc        ApplyPolicySpecsFunc
 	ApplyPolicySpecsFuncInvoked bool
 
@@ -1030,6 +1039,12 @@ type DataStore struct {
 	CleanupPolicyMembershipFunc        CleanupPolicyMembershipFunc
 	CleanupPolicyMembershipFuncInvoked bool
 
+	IncrementPolicyViolationDaysFunc        IncrementPolicyViolationDaysFunc
+	IncrementPolicyViolationDaysFuncInvoked bool
+
+	InitializePolicyViolationDaysFunc        InitializePolicyViolationDaysFunc
+	InitializePolicyViolationDaysFuncInvoked bool
+
 	LockFunc        LockFunc
 	LockFuncInvoked bool
 
@@ -1990,6 +2005,11 @@ func (s *DataStore) RecordStatisticsSent(ctx context.Context) error {
 	return s.RecordStatisticsSentFunc(ctx)
 }
 
+func (s *DataStore) CleanupStatistics(ctx context.Context) error {
+	s.CleanupStatisticsFuncInvoked = true
+	return s.CleanupStatisticsFunc(ctx)
+}
+
 func (s *DataStore) ApplyPolicySpecs(ctx context.Context, authorID uint, specs []*fleet.PolicySpec) error {
 	s.ApplyPolicySpecsFuncInvoked = true
 	return s.ApplyPolicySpecsFunc(ctx, authorID, specs)
@@ -2095,6 +2115,16 @@ func (s *DataStore) CleanupPolicyMembership(ctx context.Context, now time.Time)
 	return s.CleanupPolicyMembershipFunc(ctx, now)
 }
 
+func (s *DataStore) IncrementPolicyViolationDays(ctx context.Context) error {
+	s.IncrementPolicyViolationDaysFuncInvoked = true
+	return s.IncrementPolicyViolationDaysFunc(ctx)
+}
+
+func (s *DataStore) InitializePolicyViolationDays(ctx context.Context) error {
+	s.InitializePolicyViolationDaysFuncInvoked = true
+	return s.InitializePolicyViolationDaysFunc(ctx)
+}
+
 func (s *DataStore) Lock(ctx context.Context, name string, owner string, expiration time.Duration) (bool, error) {
 	s.LockFuncInvoked = true
 	return s.LockFunc(ctx, name, owner, expiration)