From b0c1dba44c8a8a64c593758a27b11f2aba83ec80 Mon Sep 17 00:00:00 2001 From: Juan Fernandez Date: Thu, 20 Jul 2023 08:06:43 -0400 Subject: [PATCH] Updated cache strategy on queries used in GetClientConfig (#12815) 1. Cached results of `svc.ds.Team` 2. Cached results of `svc.ds.ListQueries` too for scheduled queries only. 3. Do not load aggregated stats on `svc.ds.ListQueries` insde `GetClientConfig` --- server/datastore/cached_mysql/cached_mysql.go | 43 +++++++++ .../cached_mysql/cached_mysql_test.go | 89 +++++++++++++++++++ server/datastore/mysql/hosts_test.go | 33 +------ server/datastore/mysql/queries.go | 36 +++++++- server/datastore/mysql/queries_test.go | 65 +++++++++++++- server/datastore/mysql/teams.go | 14 +++ server/datastore/mysql/teams_test.go | 22 +++++ server/fleet/datastore.go | 5 ++ server/mock/datastore_mock.go | 24 +++++ server/service/osquery.go | 9 +- server/service/osquery_test.go | 20 +++-- 11 files changed, 312 insertions(+), 48 deletions(-) diff --git a/server/datastore/cached_mysql/cached_mysql.go b/server/datastore/cached_mysql/cached_mysql.go index 850646352f..fca6d3ed5d 100644 --- a/server/datastore/cached_mysql/cached_mysql.go +++ b/server/datastore/cached_mysql/cached_mysql.go @@ -25,6 +25,8 @@ const ( defaultTeamFeaturesExpiration = 1 * time.Minute teamMDMConfigKey = "TeamMDMConfig:team:%d" defaultTeamMDMConfigExpiration = 1 * time.Minute + teamNameByIdKey = "TeamName:team:%d" + scheduledQueriesForAgentsKey = "ScheduledQueriesAgents:team:%d" ) // cloner represents any type that can clone itself. Used by types to provide a more efficient clone method. @@ -296,10 +298,12 @@ func (ds *cachedMysql) SaveTeam(ctx context.Context, team *fleet.Team) (*fleet.T agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, team.ID) featuresKey := fmt.Sprintf(teamFeaturesKey, team.ID) mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, team.ID) + teamNameKey := fmt.Sprintf(teamNameByIdKey, team.ID) ds.c.Set(agentOptionsKey, team.Config.AgentOptions, ds.teamAgentOptionsExp) ds.c.Set(featuresKey, &team.Config.Features, ds.teamFeaturesExp) ds.c.Set(mdmConfigKey, &team.Config.MDM, ds.teamMDMConfigExp) + ds.c.Set(teamNameKey, &team.Name, ds.scheduledQueriesExp) return team, nil } @@ -313,10 +317,49 @@ func (ds *cachedMysql) DeleteTeam(ctx context.Context, teamID uint) error { agentOptionsKey := fmt.Sprintf(teamAgentOptionsKey, teamID) featuresKey := fmt.Sprintf(teamFeaturesKey, teamID) mdmConfigKey := fmt.Sprintf(teamMDMConfigKey, teamID) + teamNameKey := fmt.Sprintf(teamNameByIdKey, teamID) ds.c.Delete(agentOptionsKey) ds.c.Delete(featuresKey) ds.c.Delete(mdmConfigKey) + ds.c.Delete(teamNameKey) return nil } + +func (ds *cachedMysql) GetTeamName(ctx context.Context, teamID uint) (*string, error) { + key := fmt.Sprintf(teamNameByIdKey, teamID) + if x, found := ds.c.Get(key); found { + if teamName, ok := x.(*string); ok { + return teamName, nil + } + } + + teamName, err := ds.Datastore.GetTeamName(ctx, teamID) + if err != nil { + return nil, err + } + ds.c.Set(key, teamName, ds.scheduledQueriesExp) + return teamName, nil +} + +func (ds *cachedMysql) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + var teamIDVal uint + if teamID != nil { + teamIDVal = *teamID + } + + key := fmt.Sprintf(scheduledQueriesForAgentsKey, teamIDVal) + if x, found := ds.c.Get(key); found { + if queries, ok := x.([]*fleet.Query); ok { + return queries, nil + } + } + + queries, err := ds.Datastore.ListScheduledQueriesForAgents(ctx, teamID) + if err != nil { + return nil, err + } + ds.c.Set(key, queries, ds.scheduledQueriesExp) + return queries, nil +} diff --git a/server/datastore/cached_mysql/cached_mysql_test.go b/server/datastore/cached_mysql/cached_mysql_test.go index 0417987931..fcbde32d22 100644 --- a/server/datastore/cached_mysql/cached_mysql_test.go +++ b/server/datastore/cached_mysql/cached_mysql_test.go @@ -12,6 +12,7 @@ import ( "github.com/fleetdm/fleet/v4/server/fleet" "github.com/fleetdm/fleet/v4/server/mock" "github.com/fleetdm/fleet/v4/server/ptr" + "github.com/fleetdm/fleet/v4/server/test" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -537,3 +538,91 @@ func TestCachedTeamMDMConfig(t *testing.T) { _, err = ds.TeamMDMConfig(context.Background(), testTeam.ID) require.Error(t, err) } + +func TestCachedGetTeamName(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + mockedDS := new(mock.Store) + ds := New(mockedDS, WithScheduledQueriesExpiration(100*time.Millisecond)) + + team := fleet.Team{ + ID: 1, + CreatedAt: time.Now(), + Name: "test", + } + + deleted := false + mockedDS.GetTeamNameFunc = func(ctx context.Context, teamID uint) (*string, error) { + if deleted { + return nil, errors.New("not found") + } + return &team.Name, nil + } + mockedDS.SaveTeamFunc = func(ctx context.Context, team *fleet.Team) (*fleet.Team, error) { + return team, nil + } + mockedDS.DeleteTeamFunc = func(ctx context.Context, teamID uint) error { + deleted = true + return nil + } + + // updating updates the cache + result, err := ds.GetTeamName(ctx, 1) + require.NoError(t, err) + require.Equal(t, team.Name, *result) + + updatedTeam := &fleet.Team{ + ID: team.ID, + CreatedAt: team.CreatedAt, + Name: "test II", + } + _, err = ds.SaveTeam(ctx, updatedTeam) + require.NoError(t, err) + + result, err = ds.GetTeamName(ctx, team.ID) + require.NoError(t, err) + require.Equal(t, updatedTeam.Name, *result) + + // deleting updates the cache + err = ds.DeleteTeam(ctx, team.ID) + require.NoError(t, err) + + _, err = ds.GetTeamName(ctx, team.ID) + require.Error(t, err) +} + +func TestCachedListScheduledQueriesForAgents(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + mockedDS := new(mock.Store) + ds := New(mockedDS, WithScheduledQueriesExpiration(100*time.Millisecond)) + + teamID := ptr.Uint(1) + scheduledQueries := []*fleet.Query{ + { + ID: 1, + Name: "test", + ScheduleInterval: 100, + AutomationsEnabled: true, + TeamID: teamID, + }, + { + ID: 2, + Name: "test II", + ScheduleInterval: 100, + AutomationsEnabled: true, + TeamID: teamID, + }, + } + mockedDS.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + return scheduledQueries, nil + } + + result, err := ds.ListScheduledQueriesForAgents(ctx, teamID) + require.NoError(t, err) + test.QueryElementsMatch(t, result, scheduledQueries) +} diff --git a/server/datastore/mysql/hosts_test.go b/server/datastore/mysql/hosts_test.go index 3f84b3300d..ef2b7b3838 100644 --- a/server/datastore/mysql/hosts_test.go +++ b/server/datastore/mysql/hosts_test.go @@ -564,10 +564,6 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) { }) require.NoError(t, err) require.NoError(t, ds.AddHostsToTeam(context.Background(), &team.ID, []uint{host.ID})) - tp, err := ds.EnsureTeamPack(context.Background(), team.ID) - require.NoError(t, err) - tpQuery := test.NewQuery(t, ds, nil, "tp-time", "select * from time", 0, true) - tpSquery := test.NewScheduledQuery(t, ds, tp.ID, tpQuery.ID, 30, true, true, "time-scheduled") // Create a new pack and target to the host. // Pack and query must exist for stats to save successfully @@ -596,28 +592,8 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) { WallTime: 0, }, } - stats2 := []fleet.ScheduledQueryStats{ - { - ScheduledQueryName: tpSquery.Name, - ScheduledQueryID: tpSquery.ID, - QueryName: tpQuery.Name, - PackName: tp.Name, - PackID: tp.ID, - AverageMemory: 8000, - Denylisted: false, - Executions: 164, - Interval: 30, - LastExecuted: time.Unix(1620325191, 0).UTC(), - OutputSize: 1337, - SystemTime: 150, - UserTime: 180, - WallTime: 0, - }, - } - packStats := []fleet.PackStats{ {PackID: pack1.ID, PackName: pack1.Name, QueryStats: stats1}, - {PackID: tp.ID, PackName: teamScheduleName(team), QueryStats: stats2}, } err = ds.SaveHostPackStats(context.Background(), host.ID, packStats) require.NoError(t, err) @@ -625,14 +601,11 @@ func testHostsWithTeamPackStats(t *testing.T, ds *Datastore) { host, err = ds.Host(context.Background(), host.ID) require.NoError(t, err) - require.Len(t, host.PackStats, 2) + require.Len(t, host.PackStats, 1) sort.Sort(packStatsSlice(host.PackStats)) - assert.Equal(t, host.PackStats[0].PackName, teamScheduleName(team)) - assert.ElementsMatch(t, host.PackStats[0].QueryStats, stats2) - - assert.Equal(t, host.PackStats[1].PackName, pack1.Name) - assert.ElementsMatch(t, host.PackStats[1].QueryStats, stats1) + assert.Equal(t, host.PackStats[0].PackName, pack1.Name) + assert.ElementsMatch(t, host.PackStats[0].QueryStats, stats1) } type packStatsSlice []fleet.PackStats diff --git a/server/datastore/mysql/queries.go b/server/datastore/mysql/queries.go index 0a8af019b0..f1b9c2b48b 100644 --- a/server/datastore/mysql/queries.go +++ b/server/datastore/mysql/queries.go @@ -353,10 +353,10 @@ func (ds *Datastore) ListQueries(ctx context.Context, opt fleet.ListQueryOptions FROM queries q LEFT JOIN users u ON (q.author_id = u.id) LEFT JOIN aggregated_stats ag ON (ag.id = q.id AND ag.global_stats = ? AND ag.type = ?) - WHERE saved = true` + ` args := []interface{}{false, aggregatedStatsTypeQuery} - whereClauses := "" + whereClauses := "WHERE saved = true" if opt.OnlyObserverCanRun { whereClauses += " AND q.observer_can_run=true" @@ -453,3 +453,35 @@ func (ds *Datastore) ObserverCanRunQuery(ctx context.Context, queryID uint) (boo return observerCanRun, nil } + +func (ds *Datastore) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + sql := ` + SELECT + q.name, + q.query, + q.team_id, + q.schedule_interval, + q.platform, + q.min_osquery_version, + q.automations_enabled, + q.logging_type + FROM queries q + WHERE q.saved = true + AND (q.schedule_interval > 0 AND q.automations_enabled = 1) + ` + + args := []interface{}{} + if teamID != nil { + args = append(args, *teamID) + sql += " AND team_id = ?" + } else { + sql += " AND team_id IS NULL" + } + + results := []*fleet.Query{} + if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, sql, args...); err != nil { + return nil, ctxerr.Wrap(ctx, err, "list scheduled queries for agents") + } + + return results, nil +} diff --git a/server/datastore/mysql/queries_test.go b/server/datastore/mysql/queries_test.go index 2d82321357..3dfe9028b7 100644 --- a/server/datastore/mysql/queries_test.go +++ b/server/datastore/mysql/queries_test.go @@ -30,8 +30,9 @@ func TestQueries(t *testing.T) { {"DuplicateNew", testQueriesDuplicateNew}, {"ListFiltersObservers", testQueriesListFiltersObservers}, {"ObserverCanRunQuery", testObserverCanRunQuery}, - {"ListFiltersByTeamID", testQueriesListFiltersByTeamID}, - {"ListFiltersByIsScheduled", testQueriesListFiltersByIsScheduled}, + {"ListQueriesFiltersByTeamID", testListQueriesFiltersByTeamID}, + {"ListQueriesFiltersByIsScheduled", testListQueriesFiltersByIsScheduled}, + {"ListScheduledQueriesForAgents", testListScheduledQueriesForAgents}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -555,7 +556,7 @@ func testObserverCanRunQuery(t *testing.T, ds *Datastore) { } } -func testQueriesListFiltersByTeamID(t *testing.T, ds *Datastore) { +func testListQueriesFiltersByTeamID(t *testing.T, ds *Datastore) { globalQ1, err := ds.NewQuery(context.Background(), &fleet.Query{ Name: "query1", Query: "select 1;", @@ -617,7 +618,7 @@ func testQueriesListFiltersByTeamID(t *testing.T, ds *Datastore) { test.QueryElementsMatch(t, queries, []*fleet.Query{teamQ1, teamQ2, teamQ3}) } -func testQueriesListFiltersByIsScheduled(t *testing.T, ds *Datastore) { +func testListQueriesFiltersByIsScheduled(t *testing.T, ds *Datastore) { q1, err := ds.NewQuery(context.Background(), &fleet.Query{ Name: "query1", Query: "select 1;", @@ -669,3 +670,59 @@ func testQueriesListFiltersByIsScheduled(t *testing.T, ds *Datastore) { test.QueryElementsMatch(t, queries, tCase.expected, i) } } + +func testListScheduledQueriesForAgents(t *testing.T, ds *Datastore) { + ctx := context.Background() + + team, err := ds.NewTeam(context.Background(), &fleet.Team{ + Name: "Team 1", + Description: "Team 1", + }) + require.NoError(t, err) + + for i, teamID := range []*uint{nil, &team.ID} { + var teamIDStr string + if teamID != nil { + teamIDStr = fmt.Sprintf("%d", *teamID) + } + _, err := ds.NewQuery(context.Background(), &fleet.Query{ + Name: fmt.Sprintf("%s query1", teamIDStr), + Query: "select 1;", + Saved: true, + ScheduleInterval: 0, + TeamID: teamID, + }) + require.NoError(t, err) + _, err = ds.NewQuery(context.Background(), &fleet.Query{ + Name: fmt.Sprintf("%s query2", teamIDStr), + Query: "select 1;", + Saved: false, + ScheduleInterval: 10, + AutomationsEnabled: false, + TeamID: teamID, + }) + require.NoError(t, err) + q3, err := ds.NewQuery(context.Background(), &fleet.Query{ + Name: fmt.Sprintf("%s query3", teamIDStr), + Query: "select 1;", + Saved: true, + ScheduleInterval: 20, + AutomationsEnabled: true, + TeamID: teamID, + }) + require.NoError(t, err) + _, err = ds.NewQuery(context.Background(), &fleet.Query{ + Name: fmt.Sprintf("%s query4", teamIDStr), + Query: "select 1;", + Saved: true, + ScheduleInterval: 0, + AutomationsEnabled: true, + TeamID: teamID, + }) + require.NoError(t, err) + + result, err := ds.ListScheduledQueriesForAgents(ctx, teamID) + require.NoError(t, err) + test.QueryElementsMatch(t, result, []*fleet.Query{q3}, i) + } +} diff --git a/server/datastore/mysql/teams.go b/server/datastore/mysql/teams.go index 8aee558444..f1dd45448d 100644 --- a/server/datastore/mysql/teams.go +++ b/server/datastore/mysql/teams.go @@ -428,3 +428,17 @@ func (ds *Datastore) DeleteIntegrationsFromTeams(ctx context.Context, deletedInt } return rows.Err() } + +func (ds *Datastore) GetTeamName(ctx context.Context, teamID uint) (*string, error) { + stmt := `SELECT name FROM teams WHERE id = ?` + var teamName string + + if err := sqlx.GetContext(ctx, ds.reader(ctx), &teamName, stmt, teamID); err != nil { + if err == sql.ErrNoRows { + return nil, ctxerr.Wrap(ctx, notFound("Team").WithID(teamID)) + } + return nil, ctxerr.Wrap(ctx, err, "select team") + } + + return &teamName, nil +} diff --git a/server/datastore/mysql/teams_test.go b/server/datastore/mysql/teams_test.go index 7b46c80b51..858fa65acb 100644 --- a/server/datastore/mysql/teams_test.go +++ b/server/datastore/mysql/teams_test.go @@ -35,6 +35,7 @@ func TestTeams(t *testing.T) { {"DeleteIntegrationsFromTeams", testTeamsDeleteIntegrationsFromTeams}, {"TeamsFeatures", testTeamsFeatures}, {"TeamsMDMConfig", testTeamsMDMConfig}, + {"GetTeamByName", testGetTeamByName}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -624,3 +625,24 @@ func testTeamsMDMConfig(t *testing.T, ds *Datastore) { }, mdm) }) } + +func testGetTeamByName(t *testing.T, ds *Datastore) { + ctx := context.Background() + + t.Run("team does not exists", func(t *testing.T) { + r, err := ds.GetTeamName(ctx, 123) + require.Nil(t, r) + require.Error(t, err) + }) + + t.Run("returns the team name", func(t *testing.T) { + team, err := ds.NewTeam(ctx, &fleet.Team{ + Name: "team1", + }) + require.NoError(t, err) + + result, err := ds.GetTeamName(ctx, team.ID) + require.NoError(t, err) + require.Equal(t, team.Name, *result) + }) +} diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index 9c0792d839..85af07ab74 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -82,6 +82,9 @@ type Datastore interface { // ListQueries returns a list of queries with the provided sorting and paging options. Associated packs should also // be loaded. ListQueries(ctx context.Context, opt ListQueryOptions) ([]*Query, error) + // ListScheduledQueriesForAgents returns a list of scheduled queries (without stats) for the + // given teamID. If teamID is nil, then all scheduled queries for the 'global' team are returned. + ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*Query, error) // QueryByName looks up a query by name on a team. If teamID is nil, then the query is looked up in // the 'global' team. QueryByName(ctx context.Context, teamID *uint, name string, opts ...OptionalArg) (*Query, error) @@ -397,6 +400,8 @@ type Datastore interface { SaveTeam(ctx context.Context, team *Team) (*Team, error) // Team retrieves the Team by ID. Team(ctx context.Context, tid uint) (*Team, error) + // GetTeamName retrieves the team name by their ID. + GetTeamName(ctx context.Context, teamID uint) (*string, error) // Team deletes the Team by ID. DeleteTeam(ctx context.Context, tid uint) error // TeamByName retrieves the Team by Name. diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index e1f53f86ae..5fd0676c89 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -70,6 +70,8 @@ type QueryFunc func(ctx context.Context, id uint) (*fleet.Query, error) type ListQueriesFunc func(ctx context.Context, opt fleet.ListQueryOptions) ([]*fleet.Query, error) +type ListScheduledQueriesForAgentsFunc func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) + type QueryByNameFunc func(ctx context.Context, teamID *uint, name string, opts ...fleet.OptionalArg) (*fleet.Query, error) type ObserverCanRunQueryFunc func(ctx context.Context, queryID uint) (bool, error) @@ -300,6 +302,8 @@ type SaveTeamFunc func(ctx context.Context, team *fleet.Team) (*fleet.Team, erro type TeamFunc func(ctx context.Context, tid uint) (*fleet.Team, error) +type GetTeamNameFunc func(ctx context.Context, teamID uint) (*string, error) + type DeleteTeamFunc func(ctx context.Context, tid uint) error type TeamByNameFunc func(ctx context.Context, name string) (*fleet.Team, error) @@ -739,6 +743,9 @@ type DataStore struct { ListQueriesFunc ListQueriesFunc ListQueriesFuncInvoked bool + ListScheduledQueriesForAgentsFunc ListScheduledQueriesForAgentsFunc + ListScheduledQueriesForAgentsFuncInvoked bool + QueryByNameFunc QueryByNameFunc QueryByNameFuncInvoked bool @@ -1084,6 +1091,9 @@ type DataStore struct { TeamFunc TeamFunc TeamFuncInvoked bool + GetTeamNameFunc GetTeamNameFunc + GetTeamNameFuncInvoked bool + DeleteTeamFunc DeleteTeamFunc DeleteTeamFuncInvoked bool @@ -1809,6 +1819,13 @@ func (s *DataStore) ListQueries(ctx context.Context, opt fleet.ListQueryOptions) return s.ListQueriesFunc(ctx, opt) } +func (s *DataStore) ListScheduledQueriesForAgents(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + s.mu.Lock() + s.ListScheduledQueriesForAgentsFuncInvoked = true + s.mu.Unlock() + return s.ListScheduledQueriesForAgentsFunc(ctx, teamID) +} + func (s *DataStore) QueryByName(ctx context.Context, teamID *uint, name string, opts ...fleet.OptionalArg) (*fleet.Query, error) { s.mu.Lock() s.QueryByNameFuncInvoked = true @@ -2614,6 +2631,13 @@ func (s *DataStore) Team(ctx context.Context, tid uint) (*fleet.Team, error) { return s.TeamFunc(ctx, tid) } +func (s *DataStore) GetTeamName(ctx context.Context, teamID uint) (*string, error) { + s.mu.Lock() + s.GetTeamNameFuncInvoked = true + s.mu.Unlock() + return s.GetTeamNameFunc(ctx, teamID) +} + func (s *DataStore) DeleteTeam(ctx context.Context, tid uint) error { s.mu.Lock() s.DeleteTeamFuncInvoked = true diff --git a/server/service/osquery.go b/server/service/osquery.go index af9d8e5286..16da3b4e2f 100644 --- a/server/service/osquery.go +++ b/server/service/osquery.go @@ -347,8 +347,7 @@ func getClientConfigEndpoint(ctx context.Context, request interface{}, svc fleet } func (svc *Service) getScheduledQueries(ctx context.Context, teamID *uint) (fleet.Queries, error) { - opts := fleet.ListQueryOptions{IsScheduled: ptr.Bool(true), TeamID: teamID} - queries, err := svc.ds.ListQueries(ctx, opts) + queries, err := svc.ds.ListScheduledQueriesForAgents(ctx, teamID) if err != nil { return nil, err } @@ -444,18 +443,18 @@ func (svc *Service) GetClientConfig(ctx context.Context) (map[string]interface{} } if host.TeamID != nil { - team, err := svc.ds.Team(ctx, *host.TeamID) + teamName, err := svc.ds.GetTeamName(ctx, *host.TeamID) if err != nil { return nil, newOsqueryError("database error: " + err.Error()) } - if team != nil { + if teamName != nil { teamQueries, err := svc.getScheduledQueries(ctx, host.TeamID) if err != nil { return nil, newOsqueryError("database error: " + err.Error()) } if len(teamQueries) > 0 { - packName := fmt.Sprintf("Team: %s", team.Name) + packName := fmt.Sprintf("Team: %s", *teamName) packConfig[packName] = fleet.PackContent{ Queries: teamQueries, } diff --git a/server/service/osquery_test.go b/server/service/osquery_test.go index af76ee5aba..0b014ec45e 100644 --- a/server/service/osquery_test.go +++ b/server/service/osquery_test.go @@ -40,11 +40,9 @@ import ( func TestGetClientConfig(t *testing.T) { ds := new(mock.Store) - ds.TeamFunc = func(ctx context.Context, tid uint) (*fleet.Team, error) { - return &fleet.Team{ - Name: "Alamo", - ID: 1, - }, nil + ds.GetTeamNameFunc = func(ctx context.Context, tid uint) (*string, error) { + teamName := "Alamo" + return &teamName, nil } ds.TeamAgentOptionsFunc = func(ctx context.Context, teamID uint) (*json.RawMessage, error) { @@ -71,8 +69,8 @@ func TestGetClientConfig(t *testing.T) { return []*fleet.ScheduledQuery{}, nil } } - ds.ListQueriesFunc = func(ctx context.Context, opt fleet.ListQueryOptions) ([]*fleet.Query, error) { - if opt.TeamID == nil { + ds.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + if teamID == nil { return nil, nil } return []*fleet.Query{ @@ -2006,6 +2004,14 @@ func TestUpdateHostIntervals(t *testing.T) { svc, ctx := newTestService(t, ds, nil, nil) + ds.GetTeamNameFunc = func(ctx context.Context, tid uint) (*string, error) { + return nil, nil + } + + ds.ListScheduledQueriesForAgentsFunc = func(ctx context.Context, teamID *uint) ([]*fleet.Query, error) { + return nil, nil + } + ds.ListPacksForHostFunc = func(ctx context.Context, hid uint) ([]*fleet.Pack, error) { return []*fleet.Pack{}, nil }