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