# Copyright (C) 2010, 2011 Sebastian Thiel (byronimo@gmail.com) and contributors

#

# This module is part of GitDB and is released under

# the New BSD License: https://opensource.org/license/bsd-3-clause/

"""Test for object db"""

from gitdb.test.lib import (

TestBase,

DummyStream,

make_bytes,

make_object,

fixture_path

)

from gitdb import (

DecompressMemMapReader,

FDCompressedSha1Writer,

LooseObjectDB,

Sha1Writer,

MemoryDB,

IStream,

)

from gitdb.util import hex_to_bin

import zlib

from gitdb.typ import (

str_blob_type

)

import tempfile

import os

from io import BytesIO

class TestStream(TestBase):

"""Test stream classes"""

data_sizes = (15, 10000, 1000 * 1024 + 512)

def _assert_stream_reader(self, stream, cdata, rewind_stream=lambda s: None):

"""Make stream tests - the orig_stream is seekable, allowing it to be

rewound and reused

:param cdata: the data we expect to read from stream, the contents

:param rewind_stream: function called to rewind the stream to make it ready

for reuse"""

ns = 10

assert len(cdata) > ns - 1, "Data must be larger than %i, was %i" % (ns, len(cdata))

# read in small steps

ss = len(cdata) // ns

for i in range(ns):

data = stream.read(ss)

chunk = cdata[i * ss:(i + 1) * ss]

assert data == chunk

# END for each step

rest = stream.read()

if rest:

assert rest == cdata[-len(rest):]

# END handle rest

if isinstance(stream, DecompressMemMapReader):

assert len(stream.data()) == stream.compressed_bytes_read()

# END handle special type

rewind_stream(stream)

# read everything

rdata = stream.read()

assert rdata == cdata

if isinstance(stream, DecompressMemMapReader):

assert len(stream.data()) == stream.compressed_bytes_read()

# END handle special type

def test_decompress_reader(self):

for close_on_deletion in range(2):

for with_size in range(2):

for ds in self.data_sizes:

cdata = make_bytes(ds, randomize=False)

# zdata = zipped actual data

# cdata = original content data

# create reader

if with_size:

# need object data

zdata = zlib.compress(make_object(str_blob_type, cdata))

typ, size, reader = DecompressMemMapReader.new(zdata, close_on_deletion)

assert size == len(cdata)

assert typ == str_blob_type

# even if we don't set the size, it will be set automatically on first read

test_reader = DecompressMemMapReader(zdata, close_on_deletion=False)

assert test_reader._s == len(cdata)

else:

# here we need content data

zdata = zlib.compress(cdata)

reader = DecompressMemMapReader(zdata, close_on_deletion, len(cdata))

assert reader._s == len(cdata)

# END get reader

self._assert_stream_reader(reader, cdata, lambda r: r.seek(0))

# put in a dummy stream for closing

dummy = DummyStream()

reader._m = dummy

assert not dummy.closed

del(reader)

assert dummy.closed == close_on_deletion

# END for each datasize

# END whether size should be used

# END whether stream should be closed when deleted

def test_sha_writer(self):

writer = Sha1Writer()

assert 2 == writer.write(b"hi")

assert len(writer.sha(as_hex=1)) == 40

assert len(writer.sha(as_hex=0)) == 20

# make sure it does something ;)

prev_sha = writer.sha()

writer.write(b"hi again")

assert writer.sha() != prev_sha

def test_compressed_writer(self):

for ds in self.data_sizes:

fd, path = tempfile.mkstemp()

ostream = FDCompressedSha1Writer(fd)

data = make_bytes(ds, randomize=False)

# for now, just a single write, code doesn't care about chunking

assert len(data) == ostream.write(data)

ostream.close()

# its closed already

self.assertRaises(OSError, os.close, fd)

# read everything back, compare to data we zip

fd = os.open(path, os.O_RDONLY | getattr(os, 'O_BINARY', 0))

written_data = os.read(fd, os.path.getsize(path))

assert len(written_data) == os.path.getsize(path)

os.close(fd)

assert written_data == zlib.compress(data, 1) # best speed

os.remove(path)

# END for each os

def test_decompress_reader_special_case(self):

odb = LooseObjectDB(fixture_path('objects'))

mdb = MemoryDB()

for sha in (b'888401851f15db0eed60eb1bc29dec5ddcace911',

b'7bb839852ed5e3a069966281bb08d50012fb309b',):

ostream = odb.stream(hex_to_bin(sha))

# if there is a bug, we will be missing one byte exactly !

data = ostream.read()

assert len(data) == ostream.size

# Putting it back in should yield nothing new - after all, we have

dump = mdb.store(IStream(ostream.type, ostream.size, BytesIO(data)))

assert dump.hexsha == sha

# end for each loose object sha to test

def test_decompress_reader_chunked_read_does_not_terminate_early(self):

"""Regression test for #120: read(N) must not return b'' before EOF.

zlib can consume input without producing decompressed output (e.g.

while ingesting block headers). The reader's internal recursion

previously bailed on any empty zip output, so a caller reading in

small chunks via the standard `while chunk := stream.read(N)` idiom

would terminate at the first empty chunk -- before the actual end

of the uncompressed stream.

"""

# Highly compressible data exposes the bug because each zlib chunk

# spans many uncompressed bytes -- intermediate decompress() calls

# often return empty while consuming input.

data = b"hello world! " * 1000

zdata = zlib.compress(data)

# Loop with a small chunk size to force many internal read/decompression

# iterations before EOF.

for chunk_size in (1, 4, 16, 64):

reader = DecompressMemMapReader(

zdata, close_on_deletion=False, size=len(data)

)

out = bytearray()

while True:

chunk = reader.read(chunk_size)

if not chunk:

break

out.extend(chunk)

assert bytes(out) == data, (

f"chunk_size={chunk_size}: got {len(out)}/{len(data)} bytes"

)

assert reader._br == reader._s, (

f"chunk_size={chunk_size}: stream stopped at "

f"{reader._br}/{reader._s}"

)