utils.py: Add save_file_atomic() helper

This is a context manager that can be used to divert file writes into a
temporary file, which is then renamed into place once writing is
complete. It is primarily intended for use by source plugins which
download files, so they can ensure that their downloads appear atomic
and there is no risk of leaving half-downloaded files in the cache.

So far this is not used in the core, but it is needed by the Docker
source plugin that is proposed for the bst-external plugins repo. See:
https://gitlab.com/BuildStream/bst-external/merge_requests/9

2 files changed, 62 insertions, 0 deletions

diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
--- /dev/null
+++ b/tests/utils/__init__.py
diff --git a/tests/utils/savefile.py b/tests/utils/savefile.py
new file mode 100644
index 000000000..87f9f4b0a
--- /dev/null
+++ b/tests/utils/savefile.py
@@ -0,0 +1,62 @@
+import os
+import pytest
+
+from buildstream.utils import save_file_atomic
+
+
+def test_save_new_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-success.test')
+    with save_file_atomic(filename, 'w') as f:
+        f.write('foo\n')
+
+    assert os.listdir(tmpdir) == ['savefile-success.test']
+    with open(filename) as f:
+        assert f.read() == 'foo\n'
+
+
+def test_save_over_existing_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-overwrite.test')
+
+    with open(filename, 'w') as f:
+        f.write('existing contents\n')
+
+    with save_file_atomic(filename, 'w') as f:
+        f.write('overwritten contents\n')
+
+    assert os.listdir(tmpdir) == ['savefile-overwrite.test']
+    with open(filename) as f:
+        assert f.read() == 'overwritten contents\n'
+
+
+def test_exception_new_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-exception.test')
+
+    with pytest.raises(RuntimeError):
+        with save_file_atomic(filename, 'w') as f:
+            f.write('Some junk\n')
+            raise RuntimeError("Something goes wrong")
+
+    assert os.listdir(tmpdir) == []
+
+
+def test_exception_existing_file(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-existing.test')
+
+    with open(filename, 'w') as f:
+        f.write('existing contents\n')
+
+    with pytest.raises(RuntimeError):
+        with save_file_atomic(filename, 'w') as f:
+            f.write('Some junk\n')
+            raise RuntimeError("Something goes wrong")
+
+    assert os.listdir(tmpdir) == ['savefile-existing.test']
+    with open(filename) as f:
+        assert f.read() == 'existing contents\n'
+
+
+def test_attributes(tmpdir):
+    filename = os.path.join(tmpdir, 'savefile-attributes.test')
+    with save_file_atomic(filename, 'w') as f:
+        assert f.real_filename == filename
+        assert f.name != filename