From e0c25ba4ba9766716e5d532af730f0ac8c31a6c6 Mon Sep 17 00:00:00 2001 From: dmitri-netbird Date: Thu, 9 Jul 2026 18:17:28 +0200 Subject: [PATCH] [client] fix flaky test around event aggregation (#6710) * fix flaky test around event aggregation: control time.Now() from the test Signed-off-by: Dmitri Dolguikh * actually use passed in func to generate time Signed-off-by: Dmitri Dolguikh --------- Signed-off-by: Dmitri Dolguikh --- .../netflow/store/event_aggregation_test.go | 5 ++++- client/internal/netflow/store/memory.go | 14 ++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/client/internal/netflow/store/event_aggregation_test.go b/client/internal/netflow/store/event_aggregation_test.go index c0422e8b7..8abe0d162 100644 --- a/client/internal/netflow/store/event_aggregation_test.go +++ b/client/internal/netflow/store/event_aggregation_test.go @@ -175,7 +175,9 @@ func TestFlowAggregationOfUnknownProtocols(t *testing.T) { } func TestResetAggregationWindow(t *testing.T) { - store := NewAggregatingMemoryStore() + now := time.Now() + nowFunc := func() time.Time { return now } + store := NewAggregatingMemoryStoreWithTimeFunc(nowFunc) store.StoreEvent(&types.Event{ ID: uuid.New(), Timestamp: time.Now(), @@ -198,6 +200,7 @@ func TestResetAggregationWindow(t *testing.T) { }, }) + now = now.Add(1 * time.Second) reset := store.ResetAggregationWindow() previousEvents, ok := reset.(*AggregatingMemory) assert.True(t, ok) diff --git a/client/internal/netflow/store/memory.go b/client/internal/netflow/store/memory.go index a34e4be63..dfe764032 100644 --- a/client/internal/netflow/store/memory.go +++ b/client/internal/netflow/store/memory.go @@ -29,6 +29,7 @@ type AggregatingMemory struct { WindowStart time.Time WindowEnd time.Time rnd *v2.PCG + nowFunc func() time.Time } func (m *Memory) StoreEvent(event *types.Event) { @@ -62,14 +63,19 @@ func (m *Memory) DeleteEvents(ids []uuid.UUID) { } func NewAggregatingMemoryStore() *AggregatingMemory { - return &AggregatingMemory{WindowStart: time.Now(), Memory: Memory{events: make(map[uuid.UUID]*types.Event)}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())} + return NewAggregatingMemoryStoreWithTimeFunc(defaultNowFunc) +} + +// used in tests when deterministic (less random) time intervals are required +func NewAggregatingMemoryStoreWithTimeFunc(nowFunc func() time.Time) *AggregatingMemory { + return &AggregatingMemory{WindowStart: nowFunc(), Memory: Memory{events: make(map[uuid.UUID]*types.Event)}, nowFunc: nowFunc, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())} } func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator { am.mux.Lock() defer am.mux.Unlock() - now := time.Now() + now := am.nowFunc() toret := AggregatingMemory{WindowStart: am.WindowStart, WindowEnd: now, Memory: Memory{events: am.events}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())} am.events = make(map[uuid.UUID]*types.Event) @@ -152,3 +158,7 @@ func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event { return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here } + +func defaultNowFunc() time.Time { + return time.Now() +}