diff --git a/tests/DSUEndtoEndTest.java b/tests/DSUEndtoEndTest.java
index 8179d59..1fa661d 100644
--- a/tests/DSUEndtoEndTest.java
+++ b/tests/DSUEndtoEndTest.java
@@ -71,14 +71,25 @@
                 "Failed to fetch system image. See system_image_path parameter", imgZip);
 
         File superImg = getTempPath("super.img");
-        try (ZipFile zip = new ZipFile(imgZip)) {
-            File systemImg = getTempPath("system.img");
-            if (ZipUtil2.extractFileFromZip(zip, "system.img", systemImg)) {
+        if (imgZip.isDirectory()) {
+            File systemImg = new File(imgZip, "system.img");
+            if (systemImg.exists()) {
                 return systemImg;
             }
+            superImg = new File(imgZip, "super.img");
             Assert.assertTrue(
                     "No system.img or super.img in img zip.",
-                    ZipUtil2.extractFileFromZip(zip, "super.img", superImg));
+                    superImg.exists());
+        } else {
+            try (ZipFile zip = new ZipFile(imgZip)) {
+                File systemImg = getTempPath("system.img");
+                if (ZipUtil2.extractFileFromZip(zip, "system.img", systemImg)) {
+                    return systemImg;
+                }
+                Assert.assertTrue(
+                        "No system.img or super.img in img zip.",
+                        ZipUtil2.extractFileFromZip(zip, "super.img", superImg));
+            }
         }
 
         if (SparseImageUtil.isSparse(superImg)) {