Skip to content

Commit 479df1f

Browse files
deansheatherethanndickson
deansheather
authored and
ethanndickson
committed
Tests
1 parent c312349 commit 479df1f

File tree

3 files changed

+84
-9
lines changed

3 files changed

+84
-9
lines changed

‎coderd/coderd.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import (
3939
"cdr.dev/slog"
4040
agentproto "github.com/coder/coder/v2/agent/proto"
4141
"github.com/coder/coder/v2/buildinfo"
42+
"github.com/coder/coder/v2/clock"
4243
_ "github.com/coder/coder/v2/coderd/apidoc" // Used for swagger docs.
4344
"github.com/coder/coder/v2/coderd/appearance"
4445
"github.com/coder/coder/v2/coderd/audit"
@@ -548,6 +549,7 @@ func New(options *Options) *API {
548549
options.PrometheusRegistry.MustRegister(stn)
549550
}
550551
api.NetworkTelemetryBatcher = tailnet.NewNetworkTelemetryBatcher(
552+
clock.NewReal(),
551553
api.Options.NetworkTelemetryBatchFrequency,
552554
api.Options.NetworkTelemetryBatchMaxSize,
553555
api.handleNetworkTelemetry,

‎tailnet/service.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717

1818
"cdr.dev/slog"
1919
"github.com/coder/coder/v2/apiversion"
20+
"github.com/coder/coder/v2/clock"
2021
"github.com/coder/coder/v2/tailnet/proto"
2122
)
2223

@@ -239,19 +240,21 @@ func (c communicator) loopResp() {
239240
}
240241

241242
type NetworkTelemetryBatcher struct {
243+
clock clock.Clock
242244
frequency time.Duration
243245
maxSize int
244246
batchFn func(batch []*proto.TelemetryEvent)
245247

246248
mu sync.Mutex
247249
closed chan struct{}
248250
done chan struct{}
249-
ticker *time.Ticker
251+
ticker *clock.Ticker
250252
pending []*proto.TelemetryEvent
251253
}
252254

253-
func NewNetworkTelemetryBatcher(frequency time.Duration, maxSize int, batchFn func(batch []*proto.TelemetryEvent)) *NetworkTelemetryBatcher {
255+
func NewNetworkTelemetryBatcher(clk clock.Clock, frequency time.Duration, maxSize int, batchFn func(batch []*proto.TelemetryEvent)) *NetworkTelemetryBatcher {
254256
b := &NetworkTelemetryBatcher{
257+
clock: clk,
255258
frequency: frequency,
256259
maxSize: maxSize,
257260
batchFn: batchFn,
@@ -279,7 +282,7 @@ func (b *NetworkTelemetryBatcher) sendTelemetryBatch() {
279282
func (b *NetworkTelemetryBatcher) start() {
280283
b.mu.Lock()
281284
defer b.mu.Unlock()
282-
ticker := time.NewTicker(b.frequency)
285+
ticker := b.clock.NewTicker(b.frequency)
283286
b.ticker = ticker
284287

285288
go func() {

‎tailnet/service_test.go

Lines changed: 76 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,14 @@ import (
88
"time"
99

1010
"github.com/google/uuid"
11+
"github.com/stretchr/testify/assert"
1112
"github.com/stretchr/testify/require"
1213
"golang.org/x/xerrors"
1314
"tailscale.com/tailcfg"
1415

1516
"cdr.dev/slog"
1617
"cdr.dev/slog/sloggers/slogtest"
18+
"github.com/coder/coder/v2/clock"
1719
"github.com/coder/coder/v2/tailnet"
1820
"github.com/coder/coder/v2/tailnet/proto"
1921
"github.com/coder/coder/v2/tailnet/tailnettest"
@@ -28,12 +30,16 @@ func TestClientService_ServeClient_V2(t *testing.T) {
2830
coordPtr.Store(&coord)
2931
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
3032
derpMap := &tailcfg.DERPMap{Regions: map[int]*tailcfg.DERPRegion{999: {RegionCode: "test"}}}
33+
34+
telemetryEvents := make(chan []*proto.TelemetryEvent, 64)
3135
uut, err := tailnet.NewClientService(tailnet.ClientServiceOptions{
32-
Logger: logger,
33-
CoordPtr: &coordPtr,
34-
DERPMapUpdateFrequency: time.Millisecond,
35-
DERPMapFn: func() *tailcfg.DERPMap { return derpMap },
36-
NetworkTelemetryHandler: func(batch []*proto.TelemetryEvent) {},
36+
Logger: logger,
37+
CoordPtr: &coordPtr,
38+
DERPMapUpdateFrequency: time.Millisecond,
39+
DERPMapFn: func() *tailcfg.DERPMap { return derpMap },
40+
NetworkTelemetryHandler: func(batch []*proto.TelemetryEvent) {
41+
telemetryEvents <- batch
42+
},
3743
})
3844
require.NoError(t, err)
3945

@@ -100,7 +106,23 @@ func TestClientService_ServeClient_V2(t *testing.T) {
100106
require.NoError(t, err)
101107

102108
// PostTelemetry
103-
// TODO: write test
109+
telemetryReq := &proto.TelemetryRequest{
110+
Events: []*proto.TelemetryEvent{
111+
{
112+
Id: []byte("hi"),
113+
},
114+
{
115+
Id: []byte("bye"),
116+
},
117+
},
118+
}
119+
res, err := client.PostTelemetry(ctx, telemetryReq)
120+
require.NoError(t, err)
121+
require.NotNil(t, res)
122+
gotEvents := testutil.RequireRecvCtx(ctx, t, telemetryEvents)
123+
require.Len(t, gotEvents, 2)
124+
require.Equal(t, "hi", string(gotEvents[0].Id))
125+
require.Equal(t, "bye", string(gotEvents[1].Id))
104126

105127
// RPCs closed; we need to close the Conn to end the session.
106128
err = c.Close()
@@ -154,3 +176,51 @@ func TestClientService_ServeClient_V1(t *testing.T) {
154176
err = testutil.RequireRecvCtx(ctx, t, errCh)
155177
require.ErrorIs(t, err, expectedError)
156178
}
179+
180+
func TestNetworkTelemetryBatcher(t *testing.T) {
181+
t.Parallel()
182+
183+
var (
184+
events = make(chan []*proto.TelemetryEvent, 64)
185+
mClock = clock.NewMock(t)
186+
b = tailnet.NewNetworkTelemetryBatcher(mClock, time.Millisecond, 3, func(batch []*proto.TelemetryEvent) {
187+
assert.LessOrEqual(t, len(batch), 3)
188+
events <- batch
189+
})
190+
)
191+
192+
b.Handler([]*proto.TelemetryEvent{
193+
{Id: []byte("1")},
194+
{Id: []byte("2")},
195+
})
196+
b.Handler([]*proto.TelemetryEvent{
197+
{Id: []byte("3")},
198+
{Id: []byte("4")},
199+
})
200+
201+
// Should overflow and send a batch.
202+
ctx := testutil.Context(t, testutil.WaitShort)
203+
batch := testutil.RequireRecvCtx(ctx, t, events)
204+
require.Len(t, batch, 3)
205+
require.Equal(t, "1", string(batch[0].Id))
206+
require.Equal(t, "2", string(batch[1].Id))
207+
require.Equal(t, "3", string(batch[2].Id))
208+
209+
// Should send any pending events when the ticker fires.
210+
mClock.Advance(time.Millisecond)
211+
batch = testutil.RequireRecvCtx(ctx, t, events)
212+
require.Len(t, batch, 1)
213+
require.Equal(t, "4", string(batch[0].Id))
214+
215+
// Should send any pending events when closed.
216+
b.Handler([]*proto.TelemetryEvent{
217+
{Id: []byte("5")},
218+
{Id: []byte("6")},
219+
})
220+
err := b.Close()
221+
require.NoError(t, err)
222+
batch = testutil.RequireRecvCtx(ctx, t, events)
223+
require.Len(t, batch, 2)
224+
require.Equal(t, "5", string(batch[0].Id))
225+
require.Equal(t, "6", string(batch[1].Id))
226+
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy