Add flake support for tests

1 files changed, 61 insertions, 58 deletions

diff --git a/tests/test_snapshot.py b/tests/test_snapshot.py
index ddad410..0463418 100644
--- a/tests/test_snapshot.py
+++ b/tests/test_snapshot.py
@@ -1,74 +1,77 @@
 import unittest
+import pytest
 
 from focaccia.arch import x86
 from focaccia.snapshot import ProgramState, RegisterAccessError
 
-class TestProgramState(unittest.TestCase):
-    def setUp(self):
-        self.arch = x86.ArchX86()
+@pytest.fixture
+def arch():
+    return x86.ArchX86()
 
-    def test_register_access_empty_state(self):
-        state = ProgramState(self.arch)
-        for reg in x86.regnames:
-            self.assertRaises(RegisterAccessError, state.read_register, reg)
+@pytest.fixture
+def state(arch):
+    return ProgramState(arch)
 
-    def test_register_read_write(self):
-        state = ProgramState(self.arch)
-        for reg in x86.regnames:
-            state.set_register(reg, 0x42)
-        for reg in x86.regnames:
-            val = state.read_register(reg)
-            self.assertEqual(val, 0x42)
+@pytest.mark.parametrize("reg", x86.regnames)
+def test_register_access_empty_state(state, reg):
+    with pytest.raises(RegisterAccessError):
+        state.read_register(reg)
 
-    def test_register_aliases_empty_state(self):
-        state = ProgramState(self.arch)
-        for reg in self.arch.all_regnames:
-            self.assertRaises(RegisterAccessError, state.read_register, reg)
+def test_register_read_write(self):
+    state = ProgramState(self.arch)
+    for reg in x86.regnames:
+        state.set_register(reg, 0x42)
+    for reg in x86.regnames:
+        val = state.read_register(reg)
+        assert val == 0x42
 
-    def test_register_aliases_read_write(self):
-        state = ProgramState(self.arch)
-        for reg in ['EAX', 'EBX', 'ECX', 'EDX']:
-            state.set_register(reg, 0xa0ff0)
+def test_register_aliases_empty_state(self):
+    state = ProgramState(self.arch)
+    for reg in self.arch.all_regnames:
+        with pytest.raises(RegisterAccessError): state.read_register(reg)
 
-        for reg in ['AH', 'BH', 'CH', 'DH']:
-            self.assertEqual(state.read_register(reg), 0xf, reg)
-        for reg in ['AL', 'BL', 'CL', 'DL']:
-            self.assertEqual(state.read_register(reg), 0xf0, reg)
-        for reg in ['AX', 'BX', 'CX', 'DX']:
-            self.assertEqual(state.read_register(reg), 0x0ff0, reg)
-        for reg in ['EAX', 'EBX', 'ECX', 'EDX',
-                    'RAX', 'RBX', 'RCX', 'RDX']:
-            self.assertEqual(state.read_register(reg), 0xa0ff0, reg)
+def test_register_aliases_read_write(self):
+    state = ProgramState(self.arch)
+    for reg in ['EAX', 'EBX', 'ECX', 'EDX']:
+        state.set_register(reg, 0xa0ff0)
 
-    def test_flag_aliases(self):
-        flags = ['CF', 'PF', 'AF', 'ZF', 'SF', 'TF', 'IF', 'DF', 'OF',
-                 'IOPL', 'NT', 'RF', 'VM', 'AC', 'VIF', 'VIP', 'ID']
-        state = ProgramState(self.arch)
+    for reg in ['AH', 'BH', 'CH', 'DH']:
+        assert state.read_register(reg) == 0xf, reg
+    for reg in ['AL', 'BL', 'CL', 'DL']:
+        assert state.read_register(reg) == 0xf0, reg
+    for reg in ['AX', 'BX', 'CX', 'DX']:
+        assert state.read_register(reg) == 0x0ff0, reg
+    for reg in ['EAX', 'EBX', 'ECX', 'EDX',
+                'RAX', 'RBX', 'RCX', 'RDX']:
+        assert state.read_register(reg) == 0xa0ff0, reg
 
-        state.set_register('RFLAGS', 0)
-        for flag in flags:
-            self.assertEqual(state.read_register(flag), 0)
+def test_flag_aliases(self):
+    flags = ['CF', 'PF', 'AF', 'ZF', 'SF', 'TF', 'IF', 'DF', 'OF',
+             'IOPL', 'NT', 'RF', 'VM', 'AC', 'VIF', 'VIP', 'ID']
+    state = ProgramState(self.arch)
 
-        state.set_register('RFLAGS',
-                           x86.compose_rflags({'ZF': 1, 'PF': 1, 'OF': 0}))
-        self.assertEqual(state.read_register('ZF'), 1, self.arch.get_reg_accessor('ZF'))
-        self.assertEqual(state.read_register('PF'), 1)
-        self.assertEqual(state.read_register('OF'), 0)
-        self.assertEqual(state.read_register('AF'), 0)
-        self.assertEqual(state.read_register('ID'), 0)
-        self.assertEqual(state.read_register('SF'), 0)
+    state.set_register('RFLAGS', 0)
+    for flag in flags:
+        assert state.read_register(flag) == 0
 
-        for flag in flags:
-            state.set_register(flag, 1)
-        for flag in flags:
-            self.assertEqual(state.read_register(flag), 1)
+    state.set_register('RFLAGS',
+                       x86.compose_rflags({'ZF': 1, 'PF': 1, 'OF': 0}))
+    assert state.read_register('ZF') == 1, self.arch.get_reg_accessor('ZF')
+    assert state.read_register('PF') == 1
+    assert state.read_register('OF') == 0
+    assert state.read_register('AF') == 0
+    assert state.read_register('ID') == 0
+    assert state.read_register('SF') == 0
 
-        state.set_register('OF', 1)
-        state.set_register('AF', 1)
-        state.set_register('SF', 1)
-        self.assertEqual(state.read_register('OF'), 1)
-        self.assertEqual(state.read_register('AF'), 1)
-        self.assertEqual(state.read_register('SF'), 1)
+    for flag in flags:
+        state.set_register(flag, 1)
+    for flag in flags:
+        assert state.read_register(flag) == 1
+
+    state.set_register('OF', 1)
+    state.set_register('AF', 1)
+    state.set_register('SF', 1)
+    assert state.read_register('OF') == 1
+    assert state.read_register('AF') == 1
+    assert state.read_register('SF') == 1
 
-if __name__ == '__main__':
-    unittest.main()