Converted tests to table tests

RiskIdent · Jun 10, 2024 · 4b0156b · 4b0156b
1 parent 132cf6e
commit 4b0156b
Showing 1 changed file with 34 additions and 18 deletions.
diff --git a/plugin_test.go b/plugin_test.go
@@ -19,24 +19,40 @@ func TestInvalidConfig(t *testing.T) {
 }
 
 func TestHeaderAddress(t *testing.T) {
-	cfg := plugin.CreateConfig()
-	cfg.Headers.Address = "X-Real-Address"
-	req := testPlugin(t, cfg)
-	assertHeader(t, req.Header, "X-Real-Address", "localhost:1234")
-}
-
-func TestHeaderIP(t *testing.T) {
-	cfg := plugin.CreateConfig()
-	cfg.Headers.IP = "X-Real-IP"
-	req := testPlugin(t, cfg)
-	assertHeader(t, req.Header, "X-Real-IP", "localhost")
-}
+	tests := []struct {
+		name       string
+		cfg        plugin.ConfigHeaders
+		wantHeader string
+		wantValue  string
+	}{
+		{
+			name:       "full address",
+			cfg:        plugin.ConfigHeaders{Address: "X-Real-Address"},
+			wantHeader: "X-Real-Address",
+			wantValue:  "127.0.0.1:1234",
+		},
+		{
+			name:       "only ip",
+			cfg:        plugin.ConfigHeaders{IP: "X-Real-IP"},
+			wantHeader: "X-Real-IP",
+			wantValue:  "127.0.0.1",
+		},
+		{
+			name:       "only port",
+			cfg:        plugin.ConfigHeaders{Port: "X-Real-Port"},
+			wantHeader: "X-Real-Port",
+			wantValue:  "1234",
+		},
+	}
 
-func TestHeaderPort(t *testing.T) {
-	cfg := plugin.CreateConfig()
-	cfg.Headers.Port = "X-Real-Port"
-	req := testPlugin(t, cfg)
-	assertHeader(t, req.Header, "X-Real-Port", "1234")
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			cfg := plugin.CreateConfig()
+			cfg.Headers = tc.cfg
+			req := testPlugin(t, cfg)
+			assertHeader(t, req.Header, tc.wantHeader, tc.wantValue)
+		})
+	}
 }
 
 func testPlugin(t *testing.T, cfg *plugin.Config) *http.Request {
@@ -56,7 +72,7 @@ func testPlugin(t *testing.T, cfg *plugin.Config) *http.Request {
 		t.Fatal(err)
 	}
 
-	req.RemoteAddr = "localhost:1234"
+	req.RemoteAddr = "127.0.0.1:1234"
 	handler.ServeHTTP(recorder, req)
 
 	t.Logf("request headers: %d", len(req.Header))