diff --git a/spring-restdocs-core/src/main/java/org/springframework/restdocs/generate/RestDocumentationGenerator.java b/spring-restdocs-core/src/main/java/org/springframework/restdocs/generate/RestDocumentationGenerator.java index dcf718c6..52f9cb09 100644 --- a/spring-restdocs-core/src/main/java/org/springframework/restdocs/generate/RestDocumentationGenerator.java +++ b/spring-restdocs-core/src/main/java/org/springframework/restdocs/generate/RestDocumentationGenerator.java @@ -1,5 +1,5 @@ /* - * Copyright 2012-2016 the original author or authors. + * Copyright 2014-2016 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -62,6 +62,8 @@ public final class RestDocumentationGenerator { private final List snippets; + private final List additionalSnippets; + private final RequestConverter requestConverter; private final ResponseConverter responseConverter; @@ -167,6 +169,7 @@ public final class RestDocumentationGenerator { this.requestPreprocessor = requestPreprocessor; this.responsePreprocessor = responsePreprocessor; this.snippets = new ArrayList<>(Arrays.asList(snippets)); + this.additionalSnippets = new ArrayList<>(); } /** @@ -204,7 +207,7 @@ public final class RestDocumentationGenerator { * @param snippets the snippets to add */ public void addSnippets(Snippet... snippets) { - this.snippets.addAll(Arrays.asList(snippets)); + this.additionalSnippets.addAll(Arrays.asList(snippets)); } @SuppressWarnings("unchecked") @@ -215,6 +218,8 @@ public final class RestDocumentationGenerator { if (defaultSnippets != null) { combinedSnippets.addAll(defaultSnippets); } + combinedSnippets.addAll(this.additionalSnippets); + this.additionalSnippets.clear(); return combinedSnippets; } diff --git a/spring-restdocs-core/src/test/java/org/springframework/restdocs/RestDocumentationGeneratorTests.java b/spring-restdocs-core/src/test/java/org/springframework/restdocs/RestDocumentationGeneratorTests.java index b1887bbc..fd9b032d 100644 --- a/spring-restdocs-core/src/test/java/org/springframework/restdocs/RestDocumentationGeneratorTests.java +++ b/spring-restdocs-core/src/test/java/org/springframework/restdocs/RestDocumentationGeneratorTests.java @@ -24,6 +24,7 @@ import java.util.Map; import org.junit.Test; import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; import org.springframework.http.HttpHeaders; import org.springframework.restdocs.generate.RestDocumentationGenerator; @@ -116,18 +117,25 @@ public class RestDocumentationGeneratorTests { generator.addSnippets(additionalSnippet1, additionalSnippet2); HashMap configuration = new HashMap<>(); generator.handle(this.request, this.response, configuration); - verifySnippetInvocation(this.snippet, configuration); + generator.handle(this.request, this.response, configuration); + verifySnippetInvocation(this.snippet, configuration, 2); verifySnippetInvocation(additionalSnippet1, configuration); verifySnippetInvocation(additionalSnippet2, configuration); } private void verifySnippetInvocation(Snippet snippet, Map attributes) throws IOException { + verifySnippetInvocation(snippet, attributes, 1); + } + + private void verifySnippetInvocation(Snippet snippet, Map attributes, + int times) throws IOException { ArgumentCaptor operation = ArgumentCaptor.forClass(Operation.class); - verify(snippet).document(operation.capture()); + verify(snippet, Mockito.times(times)).document(operation.capture()); assertThat(this.operationRequest, is(equalTo(operation.getValue().getRequest()))); assertThat(this.operationResponse, is(equalTo(operation.getValue().getResponse()))); assertThat(attributes, is(equalTo(operation.getValue().getAttributes()))); } + }