Adding tests for writing

mar/mar.py

@@ -27,7 +27,7 @@ log = logging.getLogger(__name__)
 
 def read_file(fp, blocksize=8192):
     """Yields blocks of data from file object fp"""
-    for block in iter(partial(fp.read, blocksize), ''):
+    for block in iter(partial(fp.read, blocksize), b''):
         yield block
 
 
@@ -188,7 +188,7 @@ class MarInfo:
 
     def to_bytes(self):
         return struct.pack(self._member_fmt, self._offset, self.size, self.flags) + \
-            self.name + "\x00"
+            self.name.encode("ascii") + b"\x00"
 
 
 class MarFile:
@@ -236,7 +236,7 @@ class MarFile:
                 self.index_offset += 4 + 8
 
             # Write the magic and placeholder for the index
-            self.fileobj.write("MAR1" + packint(self.index_offset))
+            self.fileobj.write(b"MAR1" + packint(self.index_offset))
 
             # Write placeholder for file size
             self.fileobj.write(struct.pack(">Q", 0))

tests/test_mar.py

@@ -4,11 +4,18 @@ import os
 import tempfile
 import hashlib
 
-from mar.mar import MarFile, BZ2MarFile
+from mar.mar import MarFile, BZ2MarFile, read_file
 
 TEST_MAR = os.path.join(os.path.dirname(__file__), 'test.mar')
 
 
+def test_read_file():
+    data = []
+    for block in read_file(open(__file__, 'rb')):
+        data.append(block)
+    assert b''.join(data) == open(__file__, 'rb').read()
+
+
 def sha1sum(b):
     """Returns the sha1sum of a byte string"""
     h = hashlib.new('sha1')
@@ -26,7 +33,7 @@ def test_list():
     assert repr(m.members[1]) == "<defaults/pref/channel-prefs.js 664 76 bytes starting at 533>", m.members[1]
 
 
-class TestMar(TestCase):
+class TestReadingMar(TestCase):
     def setUp(self):
         self.tmpdir = tempfile.mkdtemp()
         self.marfile = MarFile(TEST_MAR)
@@ -48,7 +55,7 @@ class TestMar(TestCase):
         self.assertEquals("6a7890e740f1e18a425b51fefbde2f6b86f91a12", h)
 
 
-class TestBZ2Mar(TestCase):
+class TestReadingBZ2Mar(TestCase):
     def setUp(self):
         self.tmpdir = tempfile.mkdtemp()
         self.marfile = BZ2MarFile(TEST_MAR)
@@ -72,6 +79,19 @@ class TestBZ2Mar(TestCase):
         self.assertEquals("5177f5938923e94820d8565a1a0f25d19b4821d1", h)
 
 
+class TestWritingMar(TestCase):
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir)
+
+    def test_add(self):
+        marfile = os.path.join(self.tmpdir, 'test.mar')
+        with MarFile(marfile, 'w') as m:
+            m.add(__file__)
+
+
 class TestExceptions(TestCase):
     def test_badmar(self):
         self.assertRaises(ValueError, MarFile, __file__)