Tests for heartbeat

2 files changed, 115 insertions, 1 deletions

diff --git a/pkg/discord/gateway/heartbeat.go b/pkg/discord/gateway/heartbeat.go
index a8f793d..1df753a 100644
--- a/pkg/discord/gateway/heartbeat.go
+++ b/pkg/discord/gateway/heartbeat.go
@@ -50,7 +50,13 @@ func (h *HeartbeatImpl) heartbeatRoutine(interval time.Duration) {
 
 	// REF: heartbeat_interval * jitter
 	jitter := rand.Intn(int(interval))
-	time.Sleep(time.Duration(jitter))
+
+	select {
+	case <-time.After(time.Duration(jitter)):
+	case <-h.ctx.Done():
+		h.logger.Debug().Msg("heartbeat routine stopped before jitter heartbeat")
+		return
+	}
 
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
diff --git a/pkg/discord/gateway/heartbeat_test.go b/pkg/discord/gateway/heartbeat_test.go
new file mode 100644
index 0000000..f148d35
--- /dev/null
+++ b/pkg/discord/gateway/heartbeat_test.go
@@ -0,0 +1,108 @@
+package gateway_test
+
+import (
+	"context"
+	"jinx/pkg/discord/gateway"
+	"jinx/pkg/discord/mocks"
+	"testing"
+	"time"
+
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHeartBeats(t *testing.T) {
+	// This is probably a flaky test, if it fails, change the error margin.
+	// Also, Discord sets it's own interval acceptance margins, perhaps we should
+	// use that instead of a hardcoded error margin.
+	errorMargin := 1 * time.Millisecond
+
+	logger := zerolog.New(zerolog.NewTestWriter(t))
+	g := mocks.Gateway{}
+
+	beats := make(chan struct{}, 1)
+
+	g.On("Heartbeat").Return(nil).Run(func(_ mock.Arguments) {
+		beats <- struct{}{}
+	})
+
+	heartbeat := gateway.NewHeartbeat(&logger, &g)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	interval := 5 * time.Millisecond // Every 5 ms, way faster than Discord's default.
+
+	start := time.Now()
+	heartbeat.Start(ctx, interval)
+
+	<-beats
+	jitter := time.Since(start)
+	require.Less(t, jitter, interval, "the first jitter heartbeat should happen before the interval")
+	start = time.Now()
+
+	<-beats
+	secondBeat := time.Now()
+	require.WithinDuration(t, start.Add(interval), secondBeat, errorMargin, "the second beat should happen roughly with the interval")
+	start = secondBeat
+
+	<-beats
+	thirdBeat := time.Now()
+	require.WithinDuration(t, start.Add(interval), thirdBeat, errorMargin, "the third beat should happen roughly with the interval")
+
+	g.AssertNumberOfCalls(t, "Heartbeat", 3)
+	g.AssertExpectations(t)
+}
+
+func TestHeartRespondsToForce(t *testing.T) {
+	logger := zerolog.New(zerolog.NewTestWriter(t))
+	g := mocks.Gateway{}
+
+	beats := make(chan struct{}, 1)
+
+	g.On("Heartbeat").Return(nil).Run(func(_ mock.Arguments) {
+		beats <- struct{}{}
+	})
+
+	heartbeat := gateway.NewHeartbeat(&logger, &g)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	interval := 1 * time.Hour // Abnormally long, so we don't accidentally trigger the normal heartbeat.
+
+	heartbeat.Start(ctx, interval)
+
+	time.Sleep(1 * time.Millisecond) // Wait for routine to start.
+
+	start := time.Now()
+	heartbeat.ForceHeartbeat()
+	<-beats
+	require.WithinDuration(t, start, time.Now(), 1*time.Millisecond, "the forced heartbeat should happen within 1ms")
+
+	g.AssertNumberOfCalls(t, "Heartbeat", 1)
+	g.AssertExpectations(t)
+}
+
+func TestHeartDiesOnContextCancel(t *testing.T) {
+	logger := zerolog.New(zerolog.NewTestWriter(t))
+	g := mocks.Gateway{}
+
+	g.On("Heartbeat").Return(nil).Maybe()
+
+	heartbeat := gateway.NewHeartbeat(&logger, &g)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	interval := 1 * time.Millisecond
+
+	heartbeat.Start(ctx, interval)
+
+	cancel()
+
+	time.Sleep(5 * time.Millisecond)
+
+	g.AssertNotCalled(t, "Heartbeat")
+	g.AssertExpectations(t)
+}