ELFFile: allow filtering by segment type in iter_segments() (#375)

[pyelftools.git] / test / test_elffile.py

diff --git a/test/test_elffile.py b/test/test_elffile.py
index 1edba94344c710bd689458e1088dc1e487b23870..cc929b9f5beff4c7a139ae0d3fccea645d391855 100644 (file)
--- a/test/test_elffile.py
+++ b/test/test_elffile.py
@@ -14,12 +14,18 @@ class TestMap(unittest.TestCase):
     def test_address_offsets(self):
         class MockELF(ELFFile):
             __init__ = object.__init__
-            def iter_segments(self):
-                return iter((
-                    dict(p_type='PT_PHDR', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
-                    dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
-                    dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
-                ))
+            def iter_segments(self, type=None):
+                if type == 'PT_LOAD':
+                    return iter((
+                        dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                    ))
+                else:
+                    return iter((
+                        dict(p_type='PT_PHDR', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                    ))
 
         elf = MockELF()