From 8cf74ee5e3dd0a0f77be0de06157e71abdb60ad8 Mon Sep 17 00:00:00 2001
From: Simeng He <simeng@tailscale.com>
Date: Fri, 21 May 2021 12:04:04 -0400
Subject: [PATCH] Added good integration tests for PingRequest addition

---
 tstest/integration/integration_test.go        | 21 ++++++++++++++++---
 tstest/integration/testcontrol/testcontrol.go |  9 ++++----
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/tstest/integration/integration_test.go b/tstest/integration/integration_test.go
index d9e0de8fa..e45b21748 100644
--- a/tstest/integration/integration_test.go
+++ b/tstest/integration/integration_test.go
@@ -784,9 +784,24 @@ func TestControlSelectivePing(t *testing.T) {
 	n1.AwaitRunning(t)
 	n2.AwaitRunning(t)
 
-	t.Log("CONTROLURL", env.ControlServer.URL)
 	req := new(tailcfg.MapRequest)
 	req.Ping = true
-	t.Log("NEWREQ", req)
-
+	env.Control.MapResponse(req)
+	if err := tstest.WaitFor(2*time.Second, func() error {
+		st := n1.MustStatus(t)
+		req.NodeKey = tailcfg.NodeKey(st.Self.PublicKey)
+		return nil
+	}); err != nil {
+		t.Error(err)
+	}
+	mr, err := env.Control.MapResponse(req)
+	if err != nil {
+		t.Error(err)
+	}
+	t.Log("RECEIVED PR", mr.PingRequest)
+	if mr.PingRequest == nil {
+		t.Error("PingRequest does not exist")
+	}
+	d1.MustCleanShutdown(t)
+	d2.MustCleanShutdown(t)
 }
diff --git a/tstest/integration/testcontrol/testcontrol.go b/tstest/integration/testcontrol/testcontrol.go
index d3018df0d..5f135f169 100644
--- a/tstest/integration/testcontrol/testcontrol.go
+++ b/tstest/integration/testcontrol/testcontrol.go
@@ -570,11 +570,12 @@ func (s *Server) MapResponse(req *tailcfg.MapRequest) (res *tailcfg.MapResponse,
 	res.Node.AllowedIPs = res.Node.Addresses
 
 	// Function to add a PingRequest to one of its Peers to the MapResponse
-	err = s.addPingRequest(res)
-	if err != nil {
-		log.Println("ADDPINGREQ ERROR", err)
+	if req.Ping {
+		err = s.addPingRequest(res)
+		if err != nil {
+			log.Println("ADDPINGREQ ERROR", err)
+		}
 	}
-
 	return res, nil
 }