input_io_* tests

borg/testsuite/archive.py

@@ -5,6 +5,7 @@ import msgpack
 import pytest
 
 from ..archive import Archive, CacheChunkBuffer, RobustUnpacker, valid_msgpacked_dict, ITEM_KEYS
+from ..archive import InputOSError, input_io, input_io_iter
 from ..key import PlaintextKey
 from ..helpers import Manifest
 from . import BaseTestCase
@@ -145,3 +146,27 @@ def test_key_length_msgpacked_items():
     data = {key: b''}
     item_keys_serialized = [msgpack.packb(key), ]
     assert valid_msgpacked_dict(msgpack.packb(data), item_keys_serialized)
+
+
+def test_input_io():
+    with pytest.raises(InputOSError):
+        with input_io():
+            raise OSError(123)
+
+
+def test_input_io_iter():
+    class Iterator:
+        def __init__(self, exc):
+            self.exc = exc
+
+        def __next__(self):
+            raise self.exc()
+
+    oserror_iterator = Iterator(OSError)
+    with pytest.raises(InputOSError):
+        for _ in input_io_iter(oserror_iterator):
+            pass
+
+    normal_iterator = Iterator(StopIteration)
+    for _ in input_io_iter(normal_iterator):
+        assert False, 'StopIteration handled incorrectly'