add ZipPkgWriter.write()

Steve Canny · Steve Canny · commit 706464054b3b · 2013-08-05T15:08:41.000-07:00
diff --git a/opc/phys_pkg.py b/opc/phys_pkg.py
@@ -86,3 +86,10 @@ def close(self):
         releasing any resources it's using.
         """
         self._zipf.close()
+
+    def write(self, pack_uri, blob):
+        """
+        Write *blob* to this zip package with the membername corresponding to
+        *pack_uri*.
+        """
+        self._zipf.writestr(pack_uri.membername, blob)
diff --git a/tests/test_phys_pkg.py b/tests/test_phys_pkg.py
@@ -9,9 +9,14 @@
 
 """Test suite for opc.phys_pkg module."""
 
+try:
+    from io import BytesIO  # Python 3
+except ImportError:
+    from StringIO import StringIO as BytesIO
+
 import hashlib
 
-from zipfile import ZIP_DEFLATED
+from zipfile import ZIP_DEFLATED, ZipFile
 
 from opc.packuri import PACKAGE_URI, PackURI
 from opc.phys_pkg import (
@@ -110,6 +115,12 @@ def it_returns_none_when_part_has_no_rels_xml(self, phys_reader):
 
 class DescribeZipPkgWriter(object):
 
+    @pytest.fixture
+    def pkg_file(self, request):
+        pkg_file = BytesIO()
+        request.addfinalizer(pkg_file.close)
+        return pkg_file
+
     def it_opens_pkg_file_zip_on_construction(self, ZipFile_):
         pkg_file = Mock(name='pkg_file')
         ZipPkgWriter(pkg_file)
@@ -124,3 +135,19 @@ def it_can_be_closed(self, ZipFile_):
         zip_pkg_writer.close()
         # verify -----------------------
         zipf.close.assert_called_once_with()
+
+    def it_can_write_a_blob(self, pkg_file):
+        # setup ------------------------
+        pack_uri = PackURI('/part/name.xml')
+        blob = '<BlobbityFooBlob/>'.encode('utf-8')
+        # exercise ---------------------
+        pkg_writer = PhysPkgWriter(pkg_file)
+        pkg_writer.write(pack_uri, blob)
+        pkg_writer.close()
+        # verify -----------------------
+        written_blob_sha1 = hashlib.sha1(blob).hexdigest()
+        zipf = ZipFile(pkg_file, 'r')
+        retrieved_blob = zipf.read(pack_uri.membername)
+        zipf.close()
+        retrieved_blob_sha1 = hashlib.sha1(retrieved_blob).hexdigest()
+        assert retrieved_blob_sha1 == written_blob_sha1
