github.com/google/martian/v3@v3.3.3/header/header_blacklist_modifier_test.go (about)

     1  // Copyright 2015 Google Inc. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package header
    16  
    17  import (
    18  	"net/http"
    19  	"testing"
    20  
    21  	"github.com/google/martian/v3/parse"
    22  	"github.com/google/martian/v3/proxyutil"
    23  )
    24  
    25  func TestBlacklistModifierOnRequest(t *testing.T) {
    26  	mod := NewBlacklistModifier("X-Testing")
    27  
    28  	req, err := http.NewRequest("GET", "/", nil)
    29  	if err != nil {
    30  		t.Fatalf("NewRequest(): got %v, want no error", err)
    31  	}
    32  
    33  	req.Header.Set("X-Testing", "value")
    34  	req.Header.Set("Y-Testing", "value")
    35  
    36  	if err := mod.ModifyRequest(req); err != nil {
    37  		t.Fatalf("ModifyRequest(): got %v, want no error", err)
    38  	}
    39  
    40  	if _, ok := req.Header["X-Testing"]; ok {
    41  		t.Errorf("req.Header[%q]: got true, want false", "X-Testing")
    42  	}
    43  
    44  	if _, ok := req.Header["Y-Testing"]; !ok {
    45  		t.Errorf("req.Header[%q]: got false, want true", "Y-Testing")
    46  	}
    47  }
    48  
    49  func TestBlacklistModifierOnResponse(t *testing.T) {
    50  	mod := NewBlacklistModifier("X-Testing")
    51  
    52  	res := proxyutil.NewResponse(200, nil, nil)
    53  
    54  	res.Header.Set("X-Testing", "value")
    55  	res.Header.Set("Y-Testing", "value")
    56  
    57  	if err := mod.ModifyResponse(res); err != nil {
    58  		t.Fatalf("ModifyResponse(): got %v, want no error", err)
    59  	}
    60  
    61  	if _, ok := res.Header["X-Testing"]; ok {
    62  		t.Errorf("res.Header[%q]: got true, want false", "X-Testing")
    63  	}
    64  
    65  	if _, ok := res.Header["Y-Testing"]; !ok {
    66  		t.Errorf("res.Header[%q]: got false, want true", "Y-Testing")
    67  	}
    68  }
    69  
    70  func TestBlacklistModifierFromJSON(t *testing.T) {
    71  	msg := []byte(`{
    72      "header.Blacklist": {
    73    		"scope": ["request", "response"],
    74  			"names": ["X-Testing", "Y-Testing"]
    75  		}
    76  	}`)
    77  	r, err := parse.FromJSON(msg)
    78  	if err != nil {
    79  		t.Fatalf("parse.FromJSON(): got %v, want no error", err)
    80  	}
    81  
    82  	req, err := http.NewRequest("GET", "http://martian.test", nil)
    83  	if err != nil {
    84  		t.Fatalf("http.NewRequest(): got %q, want no error", err)
    85  	}
    86  
    87  	req.Header.Set("X-Testing", "value")
    88  	req.Header.Set("Y-Testing", "value")
    89  	req.Header.Set("Z-Testing", "value")
    90  
    91  	reqmod := r.RequestModifier()
    92  	if reqmod == nil {
    93  		t.Fatalf("reqmod: got nil, want not nil")
    94  	}
    95  
    96  	res := proxyutil.NewResponse(200, nil, req)
    97  	res.Header.Set("X-Testing", "value")
    98  	res.Header.Set("Y-Testing", "value")
    99  	res.Header.Set("Z-Testing", "value")
   100  
   101  	resmod := r.ResponseModifier()
   102  	if resmod == nil {
   103  		t.Fatalf("resmod: got nil, want not nil")
   104  	}
   105  
   106  	tt := []struct {
   107  		header string
   108  		want   string
   109  	}{
   110  		{
   111  			header: "X-Testing",
   112  			want:   "",
   113  		},
   114  		{
   115  			header: "Y-Testing",
   116  			want:   "",
   117  		},
   118  		{
   119  			header: "Z-Testing",
   120  			want:   "value",
   121  		},
   122  	}
   123  
   124  	for i, tc := range tt {
   125  		if err := reqmod.ModifyRequest(req); err != nil {
   126  			t.Fatalf("%d. reqmod.ModifyRequest(): got %v, want no error", i, err)
   127  		}
   128  
   129  		if got, want := req.Header.Get(tc.header), tc.want; got != want {
   130  			t.Errorf("%d. req.Header.Get(%q): got %q, want %q", i, tc.header, got, want)
   131  		}
   132  
   133  		if err := resmod.ModifyResponse(res); err != nil {
   134  			t.Fatalf("%d. resmod.ModifyResponse(): got %v, want no error", i, err)
   135  		}
   136  
   137  		if got, want := res.Header.Get(tc.header), tc.want; got != want {
   138  			t.Errorf("%d. res.Header.Get(%q): got %q, want %q", i, tc.header, got, want)
   139  		}
   140  	}
   141  }