2323import java .nio .charset .StandardCharsets ;
2424import java .nio .file .Files ;
2525import java .nio .file .Path ;
26- import java .util .ArrayList ;
2726import java .util .Collection ;
2827import java .util .Collections ;
2928import java .util .HashMap ;
29+ import java .util .HashSet ;
30+ import java .util .LinkedHashSet ;
3031import java .util .List ;
3132import java .util .Map ;
33+ import java .util .Set ;
3234import java .util .stream .Collectors ;
3335import java .util .stream .Stream ;
3436
@@ -64,7 +66,8 @@ public class WireMockContainer extends GenericContainer<WireMockContainer> {
6466
6567 private final Map <String , Stub > mappingStubs = new HashMap <>();
6668 private final Map <String , MountableFile > mappingFiles = new HashMap <>();
67- private final Map <String , Extension > extensions = new HashMap <>();
69+ private final Set <String > extensionClassNames = new LinkedHashSet <>();
70+ private final Set <File > extensionJars = new LinkedHashSet <>();
6871
6972 public WireMockContainer () {
7073 this (DEFAULT_TAG );
@@ -137,51 +140,59 @@ public WireMockContainer withFileFromResource(String name, Class<?> resource, St
137140
138141 /**
139142 * Add extension that will be loaded from the specified JAR file.
140- * @param id Unique ID of the extension, for logging purposes
143+ * @param classNames Class names of the extension to be included
144+ * @param jar JAR to be included into the container
145+ * @return this instance
146+ */
147+ public WireMockContainer withExtension (Collection <String > classNames , File jar ) {
148+ return withExtension (classNames , Collections .singleton (jar ));
149+ }
150+
151+ /**
152+ * Add extension that will be loaded from the specified JAR file.
141153 * @param classNames Class names of the extension to be included
142154 * @param jars JARs to be included into the container
143155 * @return this instance
144156 */
145- public WireMockContainer withExtension (String id , Collection <String > classNames , Collection <File > jars ) {
146- final Extension extension = new Extension (id );
147- extension .extensionClassNames .addAll (classNames );
148- extension .jars .addAll (jars );
149- extensions .put (id , extension );
157+ public WireMockContainer withExtension (Collection <String > classNames , Collection <File > jars ) {
158+ extensionClassNames .addAll (classNames );
159+ extensionJars .addAll (jars );
150160 return this ;
151161 }
152162
153163 /**
154164 * Add extension that will be loaded from the specified directory with JAR files.
155- * @param id Unique ID of the extension, for logging purposes
156165 * @param classNames Class names of the extension to be included
157- * @param jarDirectory Directory that stores all JARs
166+ * @param jarsDirectory Directory that stores all JARs
158167 * @return this instance
159168 */
160- public WireMockContainer withExtension (String id , Collection <String > classNames , File jarDirectory ) {
161- final List <File > jarsInTheDirectory ;
162- try (Stream <Path > walk = Files .walk (jarDirectory .toPath ())) {
163- jarsInTheDirectory = walk
169+ public WireMockContainer withExtension (Collection <String > classNames , Path jarsDirectory ) {
170+ if (!Files .isDirectory (jarsDirectory )) {
171+ throw new IllegalArgumentException ("Path must refers to directory " + jarsDirectory );
172+ }
173+ try (Stream <Path > walk = Files .walk (jarsDirectory )) {
174+
175+ final List <File > jarsInTheDirectory = walk
164176 .filter (p -> !Files .isDirectory (p ))
165177 .map (Path ::toFile )
166178 .filter (f -> f .toString ().endsWith (".jar" ))
167179 .collect (Collectors .toList ());
180+ return withExtension (classNames , jarsInTheDirectory );
181+
168182 } catch (IOException e ) {
169- throw new IllegalArgumentException ("Cannot list JARs in the directory " + jarDirectory , e );
183+ throw new IllegalArgumentException ("Cannot list JARs in the directory " + jarsDirectory , e );
170184 }
171-
172- return withExtension (id , classNames , jarsInTheDirectory );
173185 }
174186
175187 /**
176188 * Add extension that will be loaded from the classpath.
177189 * This method can be used if the extension is a part of the WireMock bundle,
178- * or a Jar is already added via {@link #withExtension(String, Collection, Collection)}}
179- * @param id Unique ID of the extension, for logging purposes
190+ * or a Jar is already added via {@link #withExtension(Collection, Collection)}}
180191 * @param className Class name of the extension
181192 * @return this instance
182193 */
183- public WireMockContainer withExtension (String id , String className ) {
184- return withExtension (id , Collections .singleton (className ), Collections .emptyList ());
194+ public WireMockContainer withExtension (String className ) {
195+ return withExtension (Collections .singleton (className ), Collections .emptyList ());
185196 }
186197
187198 public String getEndpoint () {
@@ -209,13 +220,8 @@ protected void configure() {
209220 withCopyToContainer (mount .getValue (), FILES_DIR + mount .getKey ());
210221 }
211222
212- final ArrayList <String > extensionClassNames = new ArrayList <>();
213- for (Map .Entry <String , Extension > entry : extensions .entrySet ()) {
214- final Extension ext = entry .getValue ();
215- extensionClassNames .addAll (ext .extensionClassNames );
216- for (File jar : ext .jars ) {
217- withCopyToContainer (MountableFile .forHostPath (jar .toPath ()), EXTENSIONS_DIR + jar .getName ());
218- }
223+ for (File jar : extensionJars ) {
224+ withCopyToContainer (MountableFile .forHostPath (jar .toPath ()), EXTENSIONS_DIR + jar .getName ());
219225 }
220226 if (!extensionClassNames .isEmpty ()) {
221227 wireMockArgs .append (" --extensions " );
@@ -236,14 +242,4 @@ public Stub (String name, String json) {
236242 }
237243 }
238244
239- private static final class Extension {
240- final String id ;
241- final List <File > jars = new ArrayList <>();
242- final List <String > extensionClassNames = new ArrayList <>();
243-
244- public Extension (String id ) {
245- this .id = id ;
246- }
247- }
248-
249245}
0 commit comments