1 files changed, 108 insertions, 0 deletions
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)
+}