Skip to content

Commit 2d5b743

Browse files
author
차지훈
committed
Provide a preprocessor that can modify existing headers (spring-projects#580)
1 parent 303656c commit 2d5b743

File tree

7 files changed

+494
-138
lines changed

7 files changed

+494
-138
lines changed

‎docs/src/docs/asciidoc/customizing-requests-and-responses.adoc

+8
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,14 @@ expression are replaced.
139139

140140

141141

142+
[[customizing-requests-and-responses-preprocessors-modify-headers]]
143+
==== Modifying Headers
144+
145+
You can use `modifyHeaders` on `Preprocessors` to add, set, and remove request
146+
or response headers.
147+
148+
149+
142150
[[customizing-requests-and-responses-preprocessors-modify-request-parameters]]
143151
==== Modifying Request Parameters
144152

‎spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessor.java

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/*
2+
* Copyright 2014-2019 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.restdocs.operation.preprocess;
18+
19+
import java.util.LinkedList;
20+
21+
import org.springframework.http.HttpHeaders;
22+
import org.springframework.restdocs.operation.OperationRequest;
23+
import org.springframework.restdocs.operation.OperationRequestFactory;
24+
import org.springframework.restdocs.operation.OperationResponse;
25+
import org.springframework.restdocs.operation.OperationResponseFactory;
26+
27+
/**
28+
* An {@link OperationPreprocessor} that can be used to modify a request's
29+
* {@link OperationRequest#getHeaders()} or response's
30+
* {@link OperationResponse#getHeaders()} by adding, setting, and removing headers.
31+
*
32+
* @author Andy Wilkinson
33+
* @author Jihun Cha
34+
*/
35+
public class HeadersModifyingOperationPreprocessor extends
36+
MultiValueMapModifyingOperationPreprocessorAdapter<HeadersModifyingOperationPreprocessor, HttpHeaders> {
37+
38+
private final OperationRequestFactory requestFactory = new OperationRequestFactory();
39+
40+
private final OperationResponseFactory responseFactory = new OperationResponseFactory();
41+
42+
@Override
43+
public OperationRequest preprocess(OperationRequest request) {
44+
HttpHeaders headers = writableHttpHeaders(request.getHeaders());
45+
for (Modification<HttpHeaders> modification : getModifications()) {
46+
modification.apply(headers);
47+
}
48+
return this.requestFactory.createFrom(request, headers);
49+
}
50+
51+
@Override
52+
public OperationResponse preprocess(OperationResponse response) {
53+
HttpHeaders headers = writableHttpHeaders(response.getHeaders());
54+
for (Modification<HttpHeaders> modification : getModifications()) {
55+
modification.apply(headers);
56+
}
57+
return this.responseFactory.createFrom(response, headers);
58+
}
59+
60+
private HttpHeaders writableHttpHeaders(HttpHeaders headers) {
61+
HttpHeaders writable = new HttpHeaders();
62+
headers.keySet()
63+
.forEach((key) -> writable.put(key, new LinkedList<>(headers.get(key))));
64+
return writable;
65+
}
66+
67+
}

‎spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/MultiValueMapModifyingOperationPreprocessorAdapter.java

+216
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
/*
2+
* Copyright 2014-2019 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.restdocs.operation.preprocess;
18+
19+
import java.util.ArrayList;
20+
import java.util.Arrays;
21+
import java.util.List;
22+
import java.util.regex.Pattern;
23+
24+
import org.springframework.util.Assert;
25+
import org.springframework.util.MultiValueMap;
26+
27+
/**
28+
* An {@link OperationPreprocessorAdapter} that supports modification of {@link T}.
29+
*
30+
* @param <S> the concrete type of this preprocessor, to be returned from methods that
31+
* support chaining
32+
* @param <T> the type to which the modification applies
33+
* @author Andy Wilkinson
34+
* @author Jihun Cha
35+
* @see ParametersModifyingOperationPreprocessor
36+
* @see HeadersModifyingOperationPreprocessor
37+
*/
38+
abstract class MultiValueMapModifyingOperationPreprocessorAdapter<S extends MultiValueMapModifyingOperationPreprocessorAdapter<S, T>, T extends MultiValueMap<String, String>>
39+
extends OperationPreprocessorAdapter {
40+
41+
private final List<Modification<T>> modifications = new ArrayList<>();
42+
43+
/**
44+
* Adds a {@link T} with the given {@code name} and {@code value}.
45+
* @param name the name
46+
* @param value the value
47+
* @return {@code this}
48+
*/
49+
@SuppressWarnings("unchecked")
50+
public S add(String name, String value) {
51+
this.modifications.add(new AddModification<>(name, value));
52+
return (S) this;
53+
}
54+
55+
/**
56+
* Sets the {@link T} with the given {@code name} to have the given {@code values}.
57+
* @param name the name
58+
* @param values the values
59+
* @return {@code this}
60+
*/
61+
@SuppressWarnings("unchecked")
62+
public S set(String name, String... values) {
63+
Assert.notEmpty(values, "At least one value must be provided");
64+
this.modifications.add(new SetModification<>(name, Arrays.asList(values)));
65+
return (S) this;
66+
}
67+
68+
/**
69+
* Removes the {@link T} with the given {@code name}.
70+
* @param name the name of the {@link T}
71+
* @return {@code this}
72+
*/
73+
@SuppressWarnings("unchecked")
74+
public S remove(String name) {
75+
this.modifications.add(new RemoveModification<>(name));
76+
return (S) this;
77+
}
78+
79+
/**
80+
* Removes the {@link T} with the given pattern of the {@code name}.
81+
* @param name the pattern of the name
82+
* @return {@code this}
83+
*/
84+
@SuppressWarnings("unchecked")
85+
public S remove(Pattern name) {
86+
this.modifications.add(new PatternRemoveModification<>(name));
87+
return (S) this;
88+
}
89+
90+
/**
91+
* Removes the given {@code value} from the {@link T} with the given {@code name}.
92+
* @param name the name
93+
* @param value the value
94+
* @return {@code this}
95+
*/
96+
@SuppressWarnings("unchecked")
97+
public S remove(String name, String value) {
98+
this.modifications.add(new RemoveValueModification<>(name, value));
99+
return (S) this;
100+
}
101+
102+
/**
103+
* Returns a list of the modifications.
104+
* @return the list of modifications
105+
*/
106+
List<Modification<T>> getModifications() {
107+
return this.modifications;
108+
}
109+
110+
/**
111+
* A {@code Modification} is used to apply the modification to {@link T}.
112+
*
113+
* @param <T> the type to which the modification applies
114+
*/
115+
interface Modification<T extends MultiValueMap<String, String>> {
116+
117+
void apply(T map);
118+
119+
}
120+
121+
private static final class AddModification<T extends MultiValueMap<String, String>>
122+
implements Modification<T> {
123+
124+
private final String name;
125+
126+
private final String value;
127+
128+
private AddModification(String name, String value) {
129+
this.name = name;
130+
this.value = value;
131+
}
132+
133+
@Override
134+
public void apply(T map) {
135+
map.add(this.name, this.value);
136+
}
137+
138+
}
139+
140+
private static final class SetModification<T extends MultiValueMap<String, String>>
141+
implements Modification<T> {
142+
143+
private final String name;
144+
145+
private final List<String> values;
146+
147+
private SetModification(String name, List<String> values) {
148+
this.name = name;
149+
this.values = values;
150+
}
151+
152+
@Override
153+
public void apply(T map) {
154+
map.put(this.name, this.values);
155+
}
156+
157+
}
158+
159+
private static final class RemoveModification<T extends MultiValueMap<String, String>>
160+
implements Modification<T> {
161+
162+
private final String name;
163+
164+
private RemoveModification(String name) {
165+
this.name = name;
166+
}
167+
168+
@Override
169+
public void apply(T map) {
170+
map.remove(this.name);
171+
}
172+
173+
}
174+
175+
private static final class PatternRemoveModification<T extends MultiValueMap<String, String>>
176+
implements Modification<T> {
177+
178+
private final Pattern name;
179+
180+
private PatternRemoveModification(Pattern name) {
181+
this.name = name;
182+
}
183+
184+
@Override
185+
public void apply(T map) {
186+
map.keySet().removeIf((key) -> this.name.matcher(key).matches());
187+
}
188+
189+
}
190+
191+
private static final class RemoveValueModification<T extends MultiValueMap<String, String>>
192+
implements Modification<T> {
193+
194+
private final String name;
195+
196+
private final String value;
197+
198+
private RemoveValueModification(String name, String value) {
199+
this.name = name;
200+
this.value = value;
201+
}
202+
203+
@Override
204+
public void apply(T map) {
205+
List<String> values = map.get(this.name);
206+
if (values != null) {
207+
values.remove(this.value);
208+
if (values.isEmpty()) {
209+
map.remove(this.name);
210+
}
211+
}
212+
}
213+
214+
}
215+
216+
}

0 commit comments

Comments
 (0)