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)