From e92b4628ff07607cef88696ddb0e3c37e2e935bd Mon Sep 17 00:00:00 2001
From: Victor Lyuboslavsky <victor@fleetdm.com>
Date: Thu, 2 May 2024 12:14:09 -0500
Subject: [PATCH] Updated datastore mock.

---
 server/mock/datastore_mock.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go
index 7f14f0d000..157695382b 100644
--- a/server/mock/datastore_mock.go
+++ b/server/mock/datastore_mock.go
@@ -435,6 +435,8 @@ type NewGlobalPolicyFunc func(ctx context.Context, authorID *uint, args fleet.Po
 
 type PolicyFunc func(ctx context.Context, id uint) (*fleet.Policy, error)
 
+type PolicyLiteFunc func(ctx context.Context, id uint) (*fleet.PolicyLite, error)
+
 type SavePolicyFunc func(ctx context.Context, p *fleet.Policy, shouldRemoveAllPolicyMemberships bool, removePolicyStats bool) error
 
 type ListGlobalPoliciesFunc func(ctx context.Context, opts fleet.ListOptions) ([]*fleet.Policy, error)
@@ -1544,6 +1546,9 @@ type DataStore struct {
 	PolicyFunc        PolicyFunc
 	PolicyFuncInvoked bool
 
+	PolicyLiteFunc        PolicyLiteFunc
+	PolicyLiteFuncInvoked bool
+
 	SavePolicyFunc        SavePolicyFunc
 	SavePolicyFuncInvoked bool
 
@@ -3729,6 +3734,13 @@ func (s *DataStore) Policy(ctx context.Context, id uint) (*fleet.Policy, error)
 	return s.PolicyFunc(ctx, id)
 }
 
+func (s *DataStore) PolicyLite(ctx context.Context, id uint) (*fleet.PolicyLite, error) {
+	s.mu.Lock()
+	s.PolicyLiteFuncInvoked = true
+	s.mu.Unlock()
+	return s.PolicyLiteFunc(ctx, id)
+}
+
 func (s *DataStore) SavePolicy(ctx context.Context, p *fleet.Policy, shouldRemoveAllPolicyMemberships bool, removePolicyStats bool) error {
 	s.mu.Lock()
 	s.SavePolicyFuncInvoked = true