Fixing live query data race in test. (#15647)

Seems like a pain to fix data races in an async test. Maybe there is an
easier way.
This commit is contained in:
Victor Lyuboslavsky 2023-12-14 14:14:43 -06:00 committed by GitHub
parent 648e292769
commit 0df5160629
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -3,6 +3,7 @@ package main
import (
"context"
"os"
"sync"
"testing"
"time"
@ -95,13 +96,22 @@ func TestLiveQuery(t *testing.T) {
ds.IsSavedQueryFunc = func(ctx context.Context, queryID uint) (bool, error) {
return true, nil
}
var GetLiveQueryStatsFuncWg sync.WaitGroup
GetLiveQueryStatsFuncWg.Add(1)
ds.GetLiveQueryStatsFunc = func(ctx context.Context, queryID uint, hostIDs []uint) ([]*fleet.LiveQueryStats, error) {
GetLiveQueryStatsFuncWg.Done()
return nil, nil
}
var UpdateLiveQueryStatsFuncWg sync.WaitGroup
UpdateLiveQueryStatsFuncWg.Add(1)
ds.UpdateLiveQueryStatsFunc = func(ctx context.Context, queryID uint, stats []*fleet.LiveQueryStats) error {
UpdateLiveQueryStatsFuncWg.Done()
return nil
}
var CalculateAggregatedPerfStatsPercentilesFuncWg sync.WaitGroup
CalculateAggregatedPerfStatsPercentilesFuncWg.Add(1)
ds.CalculateAggregatedPerfStatsPercentilesFunc = func(ctx context.Context, aggregate fleet.AggregatedStatsType, queryID uint) error {
CalculateAggregatedPerfStatsPercentilesFuncWg.Done()
return nil
}
@ -129,7 +139,21 @@ func TestLiveQuery(t *testing.T) {
expected := `{"host":"somehostname","rows":[{"bing":"fds","host_display_name":"somehostname","host_hostname":"somehostname"}]}
`
assert.Equal(t, expected, runAppForTest(t, []string{"query", "--hosts", "1234", "--query", "select 42, * from time"}))
assert.True(t, ds.GetLiveQueryStatsFuncInvoked)
assert.True(t, ds.UpdateLiveQueryStatsFuncInvoked)
assert.True(t, ds.CalculateAggregatedPerfStatsPercentilesFuncInvoked)
// We need to use waitGroups to detect whether Database functions were called because this is an asynchronous test which will flag data races otherwise.
c := make(chan struct{})
go func() {
defer close(c)
GetLiveQueryStatsFuncWg.Wait()
UpdateLiveQueryStatsFuncWg.Wait()
CalculateAggregatedPerfStatsPercentilesFuncWg.Wait()
}()
select {
case <-time.After(time.Second):
require.Fail(
t,
"Expected invocation of one of these Database functions did not happen: GetLiveQueryStats, UpdateLiveQueryStats, or CalculateAggregatedPerfStatsPercentiles",
)
case <-c: // All good
}
}