87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import os
|
|
import stat
|
|
import tempfile
|
|
import unittest
|
|
|
|
from atomic_write import atomic_write, atomic_write_text, atomic_write_bytes
|
|
|
|
|
|
class TestAtomicWrite(unittest.TestCase):
|
|
def setUp(self):
|
|
self.d = tempfile.mkdtemp()
|
|
self.p = os.path.join(self.d, "out.txt")
|
|
|
|
def test_basic_text(self):
|
|
with atomic_write(self.p) as f:
|
|
f.write("hello world")
|
|
with open(self.p) as f:
|
|
self.assertEqual(f.read(), "hello world")
|
|
|
|
def test_bytes(self):
|
|
atomic_write_bytes(self.p, b"\x00\x01\x02")
|
|
with open(self.p, "rb") as f:
|
|
self.assertEqual(f.read(), b"\x00\x01\x02")
|
|
|
|
def test_text_helper(self):
|
|
atomic_write_text(self.p, "abc")
|
|
with open(self.p) as f:
|
|
self.assertEqual(f.read(), "abc")
|
|
|
|
def test_no_partial_on_exception(self):
|
|
# pre-existing content must survive a failed write
|
|
with open(self.p, "w") as f:
|
|
f.write("ORIGINAL")
|
|
try:
|
|
with atomic_write(self.p) as f:
|
|
f.write("PARTIAL")
|
|
raise RuntimeError("boom")
|
|
except RuntimeError:
|
|
pass
|
|
with open(self.p) as f:
|
|
self.assertEqual(f.read(), "ORIGINAL")
|
|
|
|
def test_no_temp_leftover_on_exception(self):
|
|
try:
|
|
with atomic_write(self.p) as f:
|
|
f.write("x")
|
|
raise ValueError("x")
|
|
except ValueError:
|
|
pass
|
|
leftovers = [n for n in os.listdir(self.d) if n.startswith(".atomic_")]
|
|
self.assertEqual(leftovers, [])
|
|
|
|
def test_preserves_mode(self):
|
|
with open(self.p, "w") as f:
|
|
f.write("x")
|
|
os.chmod(self.p, 0o600)
|
|
with atomic_write(self.p) as f:
|
|
f.write("y")
|
|
self.assertEqual(stat.S_IMODE(os.stat(self.p).st_mode), 0o600)
|
|
|
|
def test_overwrite_replaces(self):
|
|
atomic_write_text(self.p, "first")
|
|
atomic_write_text(self.p, "second")
|
|
with open(self.p) as f:
|
|
self.assertEqual(f.read(), "second")
|
|
|
|
def test_creates_parent_dir(self):
|
|
p = os.path.join(self.d, "sub", "deep", "f.txt")
|
|
atomic_write_text(p, "nested")
|
|
with open(p) as f:
|
|
self.assertEqual(f.read(), "nested")
|
|
|
|
def test_rejects_bad_mode(self):
|
|
with self.assertRaises(ValueError):
|
|
with atomic_write(self.p, mode="a"):
|
|
pass
|
|
|
|
def test_empty_write(self):
|
|
atomic_write_text(self.p, "")
|
|
self.assertTrue(os.path.exists(self.p))
|
|
with open(self.p) as f:
|
|
self.assertEqual(f.read(), "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|