rebase proxy dialer work

Author: Divjot Arora

mongo/integration/sdam_prose_test.go

@@ -19,15 +19,13 @@ func TestSDAMProse(t *testing.T) {
 	mt := mtest.New(t)
 	defer mt.Close()
 
-	proxyDialer := newProxyDialer()
 	lowHeartbeatFrequency := 50 * time.Millisecond
 	heartbeatFrequencyClientOpts := options.Client().
-		SetHeartbeatInterval(lowHeartbeatFrequency).
-		SetDialer(proxyDialer)
+		SetHeartbeatInterval(lowHeartbeatFrequency)
 	heartbeatFrequencyMtOpts := mtest.NewOptions().
 		ClientOptions(heartbeatFrequencyClientOpts).
 		CreateCollection(false).
-		SSL(false)
+		ClientType(mtest.Proxy)
 	mt.RunOpts("heartbeats processed more frequently", heartbeatFrequencyMtOpts, func(mt *mtest.T) {
 		// Test that lowering heartbeat frequency to 50ms causes the client to process heartbeats more frequently.
 		//
@@ -47,8 +45,9 @@ func TestSDAMProse(t *testing.T) {
 		mt.Logf("num responses expected: %d\n", numExpectedResponses)
 
 		time.Sleep(timeDuration + 50*time.Millisecond)
-		assert.True(mt, len(proxyDialer.messages) >= numExpectedResponses, "expected at least %d responses, got %d",
-			numExpectedResponses, len(proxyDialer.messages))
+		messages := mt.GetProxiedMessages()
+		assert.True(mt, len(messages) >= numExpectedResponses, "expected at least %d responses, got %d",
+			numExpectedResponses, len(messages))
 	})
 
 	mt.RunOpts("rtt tests", noClientOpts, func(mt *mtest.T) {

mongo/integration/sessions_test.go

@@ -70,11 +70,15 @@ func TestSessions(t *testing.T) {
 	mtOpts := mtest.NewOptions().MinServerVersion("3.6").Topologies(mtest.ReplicaSet, mtest.Sharded).
 		CreateClient(false)
 	mt := mtest.New(t, mtOpts)
+	hosts := options.Client().ApplyURI(mt.ConnString()).Hosts
 
 	// Pin to a single mongos so heartbeats/handshakes to other mongoses won't cause errors.
+	clusterTimeClientOpts := options.Client().
+		SetHosts(hosts[:1]).
+		SetDirect(true).
+		SetHeartbeatInterval(50 * time.Second)
 	clusterTimeOpts := mtest.NewOptions().
-		ClientOptions(options.Client().SetHeartbeatInterval(50 * time.Second)).
-		ClientType(mtest.Pinned).
+		ClientOptions(clusterTimeClientOpts).
 		CreateClient(false)
 	mt.RunOpts("cluster time", clusterTimeOpts, func(mt *mtest.T) {
 		// $clusterTime included in commands
@@ -117,7 +121,6 @@ func TestSessions(t *testing.T) {
 		}
 	})
 
-	hosts := options.Client().ApplyURI(mt.ConnString()).Hosts
 	clusterTimeHandshakeOpts := options.Client().
 		SetHosts(hosts[:1]). // Prevent handshakes to other hosts from updating the cluster time.
 		SetDirect(true).
@@ -134,12 +137,12 @@ func TestSessions(t *testing.T) {
 		assert.True(mt, len(msgPairs) > 2, "expected more than two messages, got %d", len(msgPairs))
 
 		for idx, pair := range mt.GetProxiedMessages() {
-			// Get the $clusterTime value sent to the server. The first two messages are the handshakes for the
-			// heartbeat and application connections. These should not contain $clusterTime because they happen on
+			// Get the $clusterTime value sent to the server. The first three messages are the handshakes for the
+			// heartbeat, RTT, and application connections. These should not contain $clusterTime because they happen on
 			// connections that don't know the server's wire version and therefore don't know if the server supports
 			// $clusterTime.
 			_, err = pair.Sent.Command.LookupErr("$clusterTime")
-			if idx <= 1 {
+			if idx <= 2 {
 				assert.NotNil(mt, err, "expected no $clusterTime field in command %s", pair.Sent.Command)
 				continue
 			}