diff --git a/dropbox/content_hash.py b/dropbox/content_hash.py index 64078b0..bfaa7bf 100644 --- a/dropbox/content_hash.py +++ b/dropbox/content_hash.py @@ -128,8 +128,10 @@ def read(self, *args): return b def write(self, b): - self._hasher.update(b) - return self._f.write(b) + written = self._f.write(b) + if written: + self._hasher.update(b[:written]) + return written def __iter__(self): return self diff --git a/test/unit/test_content_hash.py b/test/unit/test_content_hash.py index ded83c2..ad2c42b 100644 --- a/test/unit/test_content_hash.py +++ b/test/unit/test_content_hash.py @@ -96,6 +96,32 @@ def test_write_hashes_passthrough(self): assert out.getvalue() == data assert hasher.hexdigest() == content_hash(data) + @pytest.mark.parametrize("written", [0, 3, None]) + def test_write_hashes_only_accepted_bytes(self, written): + class ShortWriter(io.BytesIO): + def write(self, data): + if written is not None: + super().write(data[:written]) + return written + + hasher = DropboxContentHasher() + out = ShortWriter() + wrapped = StreamHasher(out, hasher) + assert wrapped.write(b"partial write") == written + assert hasher.hexdigest() == content_hash(out.getvalue()) + + def test_failed_write_does_not_change_hash(self): + class FailingWriter(io.BytesIO): + def write(self, data): + raise OSError("write failed") + + hasher = DropboxContentHasher() + hasher.update(b"already written") + wrapped = StreamHasher(FailingWriter(), hasher) + with pytest.raises(OSError, match="write failed"): + wrapped.write(b"not written") + assert hasher.hexdigest() == content_hash(b"already written") + def test_readlines_hashes_and_returns_all_lines(self): lines = [b"first\n", b"second\n"] hasher = DropboxContentHasher()