community.docker/tests/unit/plugins/module_utils/test__copy.py

79 lines
1.9 KiB
Python

# Copyright 2022 Red Hat | Ansible
# GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
# SPDX-License-Identifier: GPL-3.0-or-later
from __future__ import annotations
import pytest
from ansible_collections.community.docker.plugins.module_utils._copy import (
_stream_generator_to_fileobj,
)
def _simple_generator(sequence):
yield from sequence
@pytest.mark.parametrize(
"chunks, read_sizes",
[
(
[
(1, b"1"),
(1, b"2"),
(1, b"3"),
(1, b"4"),
],
[
1,
2,
3,
],
),
(
[
(1, b"123"),
(1, b"456"),
(1, b"789"),
],
[
1,
4,
2,
2,
2,
],
),
(
[
(10 * 1024 * 1024, b"0"),
(10 * 1024 * 1024, b"1"),
],
[
1024 * 1024 - 5,
5 * 1024 * 1024 - 3,
10 * 1024 * 1024 - 2,
2 * 1024 * 1024 - 1,
2 * 1024 * 1024 + 5 + 3 + 2 + 1,
],
),
],
)
def test__stream_generator_to_fileobj(chunks, read_sizes):
chunks = [count * data for count, data in chunks]
stream = _simple_generator(chunks)
expected = b"".join(chunks)
buffer = b""
totally_read = 0
f = _stream_generator_to_fileobj(stream)
for read_size in read_sizes:
chunk = f.read(read_size)
assert len(chunk) == min(read_size, len(expected) - len(buffer))
buffer += chunk
totally_read += read_size
assert buffer == expected[: len(buffer)]
assert min(totally_read, len(expected)) == len(buffer)