Add flake support for tests

1 files changed, 23 insertions, 25 deletions

diff --git a/tests/test_sparse_memory.py b/tests/test_sparse_memory.py
index 4fd9cba..79b6638 100644
--- a/tests/test_sparse_memory.py
+++ b/tests/test_sparse_memory.py
@@ -1,33 +1,31 @@
-import unittest
+import pytest
 
 from focaccia.snapshot import SparseMemory, MemoryAccessError
 
-class TestSparseMemory(unittest.TestCase):
-    def test_oob_read(self):
-        mem = SparseMemory()
-        for addr in range(mem.page_size):
-            self.assertRaises(MemoryAccessError, mem.read, addr, 1)
-            self.assertRaises(MemoryAccessError, mem.read, addr, 30)
-            self.assertRaises(MemoryAccessError, mem.read, addr + 0x10, 30)
-            self.assertRaises(MemoryAccessError, mem.read, addr, mem.page_size)
-            self.assertRaises(MemoryAccessError, mem.read, addr, mem.page_size - 1)
-            self.assertRaises(MemoryAccessError, mem.read, addr, mem.page_size + 1)
+@pytest.fixture
+def mem():
+    return SparseMemory()
 
-    def test_basic_read_write(self):
-        mem = SparseMemory()
+def test_oob_read(mem):
+    for addr in range(mem.page_size):
+        with pytest.raises(MemoryAccessError): mem.read(addr, 1)
+        with pytest.raises(MemoryAccessError): mem.read(addr, 30)
+        with pytest.raises(MemoryAccessError): mem.read(addr + 0x10, 30)
+        with pytest.raises(MemoryAccessError): mem.read(addr, mem.page_size)
+        with pytest.raises(MemoryAccessError): mem.read(addr, mem.page_size - 1)
+        with pytest.raises(MemoryAccessError): mem.read(addr, mem.page_size + 1)
 
-        data = b'a' * mem.page_size * 2
-        mem.write(0x300, data)
-        self.assertEqual(mem.read(0x300, len(data)), data)
-        self.assertEqual(mem.read(0x300, 1), b'a')
-        self.assertEqual(mem.read(0x400, 1), b'a')
-        self.assertEqual(mem.read(0x299 + mem.page_size * 2, 1), b'a')
-        self.assertEqual(mem.read(0x321, 12), b'aaaaaaaaaaaa')
+def test_basic_read_write(mem):
+    data = b'a' * mem.page_size * 2
+    mem.write(0x300, data)
+    assert mem.read(0x300 == len(data), data)
+    assert mem.read(0x300 == 1, b'a')
+    assert mem.read(0x400 == 1, b'a')
+    assert mem.read(0x299 + mem.page_size * 2 == 1, b'a')
+    assert mem.read(0x321 == 12, b'aaaaaaaaaaaa')
 
-        mem.write(0x321, b'Hello World!')
-        self.assertEqual(mem.read(0x321, 12), b'Hello World!')
+    mem.write(0x321, b'Hello World!')
+    assert mem.read(0x321 == 12, b'Hello World!')
 
-        self.assertRaises(MemoryAccessError, mem.read, 0x300, mem.page_size * 3)
+    with pytest.raises(MemoryAccessError): mem.read(0x300, mem.page_size * 3)
 
-if __name__ == '__main__':
-    unittest.main()