UTIndex test case updated into proper one

construct · Jan 30, 2021 · 0c77120 · 0c77120
1 parent 65d40e5
commit 0c77120
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 41 deletions.
diff --git a/gallery/__init__.py b/gallery/__init__.py
@@ -1 +1,2 @@
 from .pe32coff import pe32file
+from .ut_index import UTIndex
diff --git a/gallery/ut_index.py b/gallery/ut_index.py
@@ -1,18 +1,8 @@
-import sys
-sys.path = [".."] + sys.path  # Use construct from repo, not pip
-from construct import (
-    Construct,
-    byte2int,
-    int2byte,
-    stream_read,
-    stream_write,
-    integertypes,
-    IntegerError,
-    Struct,
-)
+from construct import *
+from construct.lib import *
 
 
-class Idx(Construct):
+class UTIndex(Construct):
     """
     Format for "Index" objects in Unreal Tournament 1999 packages.
     Index objects are variable length signed integers with the following structure:
@@ -80,30 +70,3 @@ def _build(self, obj, stream, context, path):
             if not more_bit:
                 break
         return obj
-
-
-if __name__ == "__main__":
-    c = Struct("index" / Idx())
-    test_data = [
-        [0x0f, 0x40, 0xff],  # 0x0f
-        [0x4f, 0x40, 0xff],  # (0x40 << 6) + 0x0f = 0x100f
-        [0x8f, 0x40, 0xff],  # -0x0f
-        [0xcf, 0x40, 0xff],  # -((0x40 << 6) + 0x0f) = 0x100f
-        [0x4f, 0x80, 0x40, 0xff],  # (0x40 << 13) + 0x0f = 0x8000f
-        [0x4f, 0x80, 0x80, 0x40, 0xff],  # (0x40 << 20) + 0x0f = 0x400000f
-        [0x4f, 0x80, 0x80, 0x80, 0x8f, 0xff]  # 0x8f << 27 + 0x0f = 0x47800000f
-    ]
-    expected_values = [
-        0x0f,
-        0x100f,
-        -0x0f,
-        -0x100f,
-        0x8000f,
-        0x400000f,
-        0x47800000f
-    ]
-    for test, ev in zip(test_data, expected_values):
-        assert c.parse(bytes(test)).index == ev
-        new_bytes = c.build(dict(index=ev))
-        assert new_bytes == bytes(test[:len(new_bytes)])
-    print("All tests passed!")
diff --git a/tests/gallery/test_gallery.py b/tests/gallery/test_gallery.py
@@ -2,11 +2,37 @@
 from construct import *
 from construct.lib import *
 
-from gallery import pe32file
+from gallery import pe32file, UTIndex
 
 
 def test_pe32():
     commondump(pe32file, "python37-win32.exe")
     commondump(pe32file, "python37-win64.exe")
     commondump(pe32file, "SharpZipLib0860-dotnet20.dll")
     commondump(pe32file, "sqlite3.dll")
+
+def test_utindex():
+    d = UTIndex()
+    test_data = [
+        [0x0f, 0x40, 0xff],  # 0x0f
+        [0x4f, 0x40, 0xff],  # (0x40 << 6) + 0x0f = 0x100f
+        [0x8f, 0x40, 0xff],  # -0x0f
+        [0xcf, 0x40, 0xff],  # -((0x40 << 6) + 0x0f) = 0x100f
+        [0x4f, 0x80, 0x40, 0xff],  # (0x40 << 13) + 0x0f = 0x8000f
+        [0x4f, 0x80, 0x80, 0x40, 0xff],  # (0x40 << 20) + 0x0f = 0x400000f
+        [0x4f, 0x80, 0x80, 0x80, 0x8f, 0xff]  # 0x8f << 27 + 0x0f = 0x47800000f
+    ]
+    expected_values = [
+        0x0f,
+        0x100f,
+        -0x0f,
+        -0x100f,
+        0x8000f,
+        0x400000f,
+        0x47800000f
+    ]
+    for test, ev in zip(test_data, expected_values):
+        assert d.parse(bytes(test)) == ev
+        new_bytes = d.build(ev)
+        assert new_bytes == bytes(test[:len(new_bytes)])
+    print("All tests passed!")