From 488af3dd4ae3ade7dc86555b542a7f098afd79d7 Mon Sep 17 00:00:00 2001
From: Qingyang Hu <103950869+qingyang-hu@users.noreply.github.com>
Date: Tue, 6 Feb 2024 20:40:21 -0500
Subject: [PATCH] GODRIVER-3036 Add prose test for ServerHeartbeatStartedEvent.
 (#1548)

---
 mongo/integration/sdam_prose_test.go | 65 ++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/mongo/integration/sdam_prose_test.go b/mongo/integration/sdam_prose_test.go
index 435bdc72da..f91bab1176 100644
--- a/mongo/integration/sdam_prose_test.go
+++ b/mongo/integration/sdam_prose_test.go
@@ -8,16 +8,22 @@ package integration
 
 import (
 	"context"
+	"net"
 	"os"
 	"runtime"
 	"testing"
 	"time"
 
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/event"
 	"go.mongodb.org/mongo-driver/internal/assert"
 	"go.mongodb.org/mongo-driver/internal/handshake"
+	"go.mongodb.org/mongo-driver/internal/require"
+	"go.mongodb.org/mongo-driver/mongo/address"
 	"go.mongodb.org/mongo-driver/mongo/description"
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/topology"
 )
 
 func TestSDAMProse(t *testing.T) {
@@ -173,3 +179,62 @@ func TestSDAMProse(t *testing.T) {
 
 	})
 }
+
+func TestServerHeartbeatStartedEvent(t *testing.T) {
+	t.Run("emits the first HeartbeatStartedEvent before the monitoring socket was created", func(t *testing.T) {
+		t.Parallel()
+
+		const address = address.Address("localhost:9999")
+		expectedEvents := []string{
+			"serverHeartbeatStartedEvent",
+			"client connected",
+			"client hello received",
+			"serverHeartbeatFailedEvent",
+		}
+
+		events := make(chan string)
+
+		listener, err := net.Listen("tcp", address.String())
+		assert.NoError(t, err)
+		defer listener.Close()
+		go func() {
+			conn, err := listener.Accept()
+			assert.NoError(t, err)
+			defer conn.Close()
+
+			events <- "client connected"
+			_, _ = conn.Read(nil)
+			events <- "client hello received"
+		}()
+
+		server := topology.NewServer(
+			address,
+			primitive.NewObjectID(),
+			topology.WithServerMonitor(func(*event.ServerMonitor) *event.ServerMonitor {
+				return &event.ServerMonitor{
+					ServerHeartbeatStarted: func(e *event.ServerHeartbeatStartedEvent) {
+						events <- "serverHeartbeatStartedEvent"
+					},
+					ServerHeartbeatFailed: func(e *event.ServerHeartbeatFailedEvent) {
+						events <- "serverHeartbeatFailedEvent"
+					},
+				}
+			}),
+		)
+		require.NoError(t, server.Connect(nil))
+
+		ticker := time.NewTicker(5 * time.Second)
+		defer ticker.Stop()
+
+		actualEvents := make([]string, 0, len(expectedEvents))
+		for len(actualEvents) < len(expectedEvents) {
+			select {
+			case event := <-events:
+				actualEvents = append(actualEvents, event)
+			case <-ticker.C:
+				assert.FailNow(t, "timed out for incoming event")
+			}
+		}
+		assert.Equal(t, expectedEvents, actualEvents)
+	})
+}