adding tests and delete behavior

Signed-off-by: Jesse Jaggars <jjaggars@redhat.com>

Author: jhjaggars

receptor/buffers/file.py

@@ -39,11 +39,15 @@ async def put(self, data):
             await self.q.put(ident)
             await self._save_manifest()
 
-    async def get(self, handle_only=False):
+    async def get(self, handle_only=False, delete=True):
         async with self._manifest_lock:
-            msg = await self.q.get()
-            await self._save_manifest()
-            return await self._get_file(msg, handle_only=handle_only)
+            while True:
+                msg = await self.q.get()
+                await self._save_manifest()
+                try:
+                    return await self._get_file(msg, handle_only=handle_only, delete=delete)
+                except FileNotFoundError:
+                    pass
 
     async def _save_manifest(self):
         await self._loop.run_in_executor(pool, self._write_manifest)
@@ -59,13 +63,22 @@ def _read_manifest(self):
         except FileNotFoundError:
             return []
 
-    async def _get_file(self, path, handle_only=False):
+    async def _get_file(self, path, handle_only=False, delete=True):
+        """
+        Retrieves a file from disk. If handle_only is True then we will
+        return the handle to the file and do nothing else. Otherwise the file
+        is read into memory all at once and returned. If delete is True (the
+        default) and handle_only is False (the default) then the underlying
+        file will be removed as well.
+        """
         path = os.path.join(self._message_path, path)
         fp = await self._loop.run_in_executor(pool, open, path, "rb")
         if handle_only:
             return fp
         bytes = await self._loop.run_in_executor(pool, lambda: fp.read())
         fp.close()
+        if delete:
+            os.remove(path)
         return bytes
 
     def _write_file(self, data, ident):

receptor/tests/test_durable_buffer.py

@@ -0,0 +1,52 @@
+import os
+import shutil
+import tempfile
+
+import pytest
+
+from receptor.buffers.file import DurableBuffer
+
+
+@pytest.yield_fixture
+def tempdir():
+    dir_ = tempfile.mkdtemp()
+    yield dir_
+    shutil.rmtree(dir_)
+
+
+@pytest.mark.asyncio
+async def test_create(event_loop, tempdir):
+    b = DurableBuffer(tempdir, "test_create", event_loop)
+    await b.put(b"some data")
+    assert await b.get() == b"some data"
+
+
+@pytest.mark.asyncio
+async def test_manifest(event_loop, tempdir):
+    b = DurableBuffer(tempdir, "test_manifest", event_loop)
+    await b.put(b"one")
+    await b.put(b"two")
+    await b.put(b"three")
+    assert b.q.qsize() == 3
+
+    assert len(b._read_manifest()) == 3
+
+
+@pytest.mark.asyncio
+async def test_unreadable_file(event_loop, tempdir):
+    b = DurableBuffer(tempdir, "test_unreadable_file", event_loop)
+    b.q._queue.appendleft("junk")
+    await b.put(b"valid data")
+    data = await b.get()
+    assert data == b"valid data"
+    assert b.q.empty()
+
+
+@pytest.mark.asyncio
+async def test_deletes_messages(event_loop, tempdir):
+    b = DurableBuffer(tempdir, "test_deletes_messages", event_loop)
+    await b.put(b"some data")
+    ident = b.q._queue[0]
+    assert await b.get() == b"some data"
+    filepath = os.path.join(b._message_path, ident)
+    assert not os.path.exists(filepath)