From 10662c4282aa7c137ec582f24ab98116fb0de14e Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 12 Aug 2024 20:36:31 -0700
Subject: [PATCH] tstest/integration/nat: annotate test 'want' values, fail on
 mismatch

Updates #13038

Change-Id: Id711ee19e52a7051a2273c806b184c5571c6e24f
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tstest/integration/nat/nat_test.go | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/tstest/integration/nat/nat_test.go b/tstest/integration/nat/nat_test.go
index 51fa0abbd..788a47615 100644
--- a/tstest/integration/nat/nat_test.go
+++ b/tstest/integration/nat/nat_test.go
@@ -44,6 +44,8 @@ type natTest struct {
 	tempDir string // for qcow2 images
 	vnet    *vnet.Server
 	kernel  string // linux kernel path
+
+	gotRoute pingRoute
 }
 
 func newNatTest(tb testing.TB) *natTest {
@@ -328,9 +330,10 @@ func (nt *natTest) runTest(node1, node2 addNodeFunc) pingRoute {
 	if err != nil {
 		t.Fatalf("ping failure: %v", err)
 	}
-	route := classifyPing(pingRes)
-	t.Logf("ping route: %v", route)
-	return route
+	nt.gotRoute = classifyPing(pingRes)
+	t.Logf("ping route: %v", nt.gotRoute)
+
+	return nt.gotRoute
 }
 
 func classifyPing(pr *ipnstate.PingResult) pingRoute {
@@ -450,34 +453,51 @@ var types = []nodeType{
 	{"sameLAN", sameLAN},
 }
 
+// want sets the expected ping route for the test.
+func (nt *natTest) want(r pingRoute) {
+	if nt.gotRoute != r {
+		nt.tb.Errorf("ping route = %v; want %v", nt.gotRoute, r)
+	}
+}
+
 func TestEasyEasy(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easy, easy)
+	nt.want(routeDirect)
 }
 
-// Tests https://github.com/tailscale/tailscale/issues/3824 ...
+func TestSameLAN(t *testing.T) {
+	nt := newNatTest(t)
+	nt.runTest(easy, sameLAN)
+	nt.want(routeLocal)
+}
+
+// TestBPFDisco tests https://github.com/tailscale/tailscale/issues/3824 ...
 // * server behind a Hard NAT
 // * client behind a NAT with UPnP support
 // * client machine has a stateful host firewall (e.g. ufw)
-
 func TestBPFDisco(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easyPMPFW, hard)
+	nt.want(routeDirect)
 }
 
 func TestHostFWNoBPF(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easyPMPFWNoBPF, hard)
+	nt.want(routeDERP)
 }
 
 func TestHostFWPair(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easyFW, easyFW)
+	nt.want(routeDirect)
 }
 
 func TestOneHostFW(t *testing.T) {
 	nt := newNatTest(t)
 	nt.runTest(easy, easyFW)
+	nt.want(routeDirect)
 }
 
 var pair = flag.String("pair", "", "comma-separated pair of types to test (easy, easyAF, hard, easyPMP, hardPMP, one2one, sameLAN)")