Files
dvc/tests/unit/utils/test_stream.py
T

75 lines
2.2 KiB
Python

import pytest
from dvc.fs.local import LocalFileSystem
from dvc.istextfile import DEFAULT_CHUNK_SIZE, istextfile
from dvc.utils import file_md5
from dvc.utils.stream import HashedStreamReader
def test_hashed_stream_reader(tmp_dir):
tmp_dir.gen({"foo": "foo"})
foo = tmp_dir / "foo"
with open(foo, "rb") as fobj:
stream_reader = HashedStreamReader(fobj)
assert stream_reader.readable()
assert not stream_reader.seekable()
assert stream_reader.read(2) == b"fo"
assert stream_reader.tell() == 2
assert stream_reader.read(1) == b"o"
assert stream_reader.tell() == 3
hex_digest = file_md5(foo, LocalFileSystem())
assert stream_reader.is_text_file
assert hex_digest == stream_reader.hash_info.value
def test_hashed_stream_reader_as_chunks(tmp_dir):
tmp_dir.gen({"foo": b"foo \x00" * 16})
foo = tmp_dir / "foo"
actual_size = len(foo.read_bytes())
with open(foo, "rb") as fobj:
stream_reader = HashedStreamReader(fobj)
total_read = 0
while True:
chunk = stream_reader.read(16)
total_read += len(chunk)
assert stream_reader.tell() == total_read
if not chunk:
break
assert stream_reader.tell() == actual_size == total_read
hex_digest = file_md5(foo, LocalFileSystem())
assert not stream_reader.is_text_file
assert hex_digest == stream_reader.hash_info.value
@pytest.mark.parametrize(
"contents",
[b"x" * DEFAULT_CHUNK_SIZE + b"\x00", b"clean", b"not clean \x00"],
)
def test_hashed_stream_reader_compatibility(tmp_dir, contents):
# Always read more than the DEFAULT_CHUNK_SIZE (512 bytes).
# This imitates the read actions performed by upload_fobj.
chunk_size = DEFAULT_CHUNK_SIZE * 2
tmp_dir.gen("data", contents)
data = tmp_dir / "data"
with open(data, "rb") as fobj:
stream_reader = HashedStreamReader(fobj)
stream_reader.read(chunk_size)
local_fs = LocalFileSystem()
hex_digest = file_md5(data, local_fs)
assert stream_reader.is_text_file is istextfile(data, local_fs)
assert stream_reader.hash_info.value == hex_digest