Cleanup (test)

1 files changed, 23 insertions, 8 deletions

diff --git a/test/org/jetbrains/java/decompiler/SingleClassesTest.java b/test/org/jetbrains/java/decompiler/SingleClassesTest.java
index 9868011..cdd707f 100644
--- a/test/org/jetbrains/java/decompiler/SingleClassesTest.java
+++ b/test/org/jetbrains/java/decompiler/SingleClassesTest.java
@@ -21,6 +21,9 @@ import org.junit.Before;
 import org.junit.Test;
 
 import java.io.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -60,14 +63,7 @@ public class SingleClassesTest {
       assertTrue(classFile.isFile());
 
       ConsoleDecompiler decompiler = fixture.getDecompiler();
-      decompiler.addSpace(classFile, true);
-      File[] innerClasses = classFile.getParentFile().listFiles(new FilenameFilter() {
-        @Override
-        public boolean accept(File dir, String name) {
-          return name.matches(testName + "\\$.+\\.class");
-        }
-      });
-      for (File inner : innerClasses) {
+      for (File inner : collectClasses(classFile)) {
         decompiler.addSpace(inner, true);
       }
 
@@ -88,6 +84,25 @@ public class SingleClassesTest {
     }
   }
 
+  private static List<File> collectClasses(File classFile) {
+    List<File> files = new ArrayList<File>();
+    files.add(classFile);
+
+    File parent = classFile.getParentFile();
+    if (parent != null) {
+      final String pattern = classFile.getName().replace(".class", "") + "\\$.+\\.class";
+      File[] inner = parent.listFiles(new FilenameFilter() {
+        @Override
+        public boolean accept(File dir, String name) {
+          return name.matches(pattern);
+        }
+      });
+      if (inner != null) Collections.addAll(files, inner);
+    }
+
+    return files;
+  }
+
   private static String getContent(File file) throws IOException {
     Reader reader = new InputStreamReader(new FileInputStream(file), "UTF-8");
     try {