From 5f45d8f8e66d581d5ad9d25d1e2fe606a4fac156 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 10 Aug 2021 13:38:12 -0700
Subject: [PATCH] tsweb: add VarzHandler tests

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tsweb/tsweb.go      |   4 +-
 tsweb/tsweb_test.go | 105 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 6a467a487..ab2cfe924 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -424,11 +424,13 @@ func VarzHandler(w http.ResponseWriter, r *http.Request) {
 			})
 		}
 	}
-	expvar.Do(func(kv expvar.KeyValue) {
+	expvarDo(func(kv expvar.KeyValue) {
 		dump("", kv)
 	})
 }
 
+var expvarDo = expvar.Do // pulled out for tests
+
 func writeMemstats(w io.Writer, ms *runtime.MemStats) {
 	out := func(name, typ string, v uint64, help string) {
 		if help != "" {
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index 5e0dfe9c2..f9bda99cc 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"context"
 	"errors"
+	"expvar"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -15,6 +16,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"tailscale.com/metrics"
 	"tailscale.com/tstest"
 )
 
@@ -300,3 +302,106 @@ func BenchmarkLog(b *testing.B) {
 		h.ServeHTTP(rw, req)
 	}
 }
+
+func TestVarzHandler(t *testing.T) {
+	tests := []struct {
+		name string
+		k    string // key name
+		v    expvar.Var
+		want string
+	}{
+		{
+			"int",
+			"foo",
+			new(expvar.Int),
+			"# TYPE foo counter\nfoo 0\n",
+		},
+		{
+			"int_with_type_counter",
+			"counter_foo",
+			new(expvar.Int),
+			"# TYPE foo counter\nfoo 0\n",
+		},
+		{
+			"int_with_type_gauge",
+			"gauge_foo",
+			new(expvar.Int),
+			"# TYPE foo gauge\nfoo 0\n",
+		},
+		{
+			"metrics_set",
+			"s",
+			&metrics.Set{
+				Map: *(func() *expvar.Map {
+					m := new(expvar.Map)
+					m.Init()
+					m.Add("foo", 1)
+					m.Add("bar", 2)
+					return m
+				})(),
+			},
+			"# TYPE s_bar counter\ns_bar 2\n# TYPE s_foo counter\ns_foo 1\n",
+		},
+		{
+			"metrics_set_TODO_guage_type",
+			"gauge_s", // TODO(bradfitz): arguably a bug; should pass down type
+			&metrics.Set{
+				Map: *(func() *expvar.Map {
+					m := new(expvar.Map)
+					m.Init()
+					m.Add("foo", 1)
+					m.Add("bar", 2)
+					return m
+				})(),
+			},
+			"# TYPE s_bar counter\ns_bar 2\n# TYPE s_foo counter\ns_foo 1\n",
+		},
+		{
+			"func_float64",
+			"counter_x",
+			expvar.Func(func() interface{} { return float64(1.2) }),
+			"# TYPE x counter\nx 1.2\n",
+		},
+		{
+			"func_float64_gauge",
+			"gauge_x",
+			expvar.Func(func() interface{} { return float64(1.2) }),
+			"# TYPE x gauge\nx 1.2\n",
+		},
+		{
+			"func_float64_untyped",
+			"x",
+			expvar.Func(func() interface{} { return float64(1.2) }),
+			"# skipping expvar \"x\" (Go type expvar.Func returning float64) with undeclared Prometheus type\n",
+		},
+		{
+			"label_map",
+			"counter_m",
+			&metrics.LabelMap{
+				Label: "label",
+				Map: *(func() *expvar.Map {
+					m := new(expvar.Map)
+					m.Init()
+					m.Add("foo", 1)
+					m.Add("bar", 2)
+					return m
+				})(),
+			},
+			"# TYPE m counter\nm{label=\"bar\"} 2\nm{label=\"foo\"} 1\n",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer func() { expvarDo = expvar.Do }()
+			expvarDo = func(f func(expvar.KeyValue)) {
+				f(expvar.KeyValue{Key: tt.k, Value: tt.v})
+			}
+			rec := httptest.NewRecorder()
+			VarzHandler(rec, httptest.NewRequest("GET", "/", nil))
+			if got := rec.Body.Bytes(); string(got) != tt.want {
+				t.Errorf("mismatch\n got: %q\nwant: %q\n", got, tt.want)
+			}
+		})
+	}
+
+}