diff --git a/aiocogeo/tag.py b/aiocogeo/tag.py index 9fc10f8..e320cdc 100644 --- a/aiocogeo/tag.py +++ b/aiocogeo/tag.py @@ -4,13 +4,13 @@ from typing import Any, Optional, Tuple, Union -from .config import HEADER_CHUNK_SIZE, LOG_LEVEL +from . import config from .constants import GEO_KEYS, TIFF_TAGS from .filesystems import Filesystem logger = logging.getLogger(__name__) -logger.setLevel(LOG_LEVEL) +logger.setLevel(config.LOG_LEVEL) @dataclass @@ -89,7 +89,7 @@ async def read(cls, reader: Filesystem) -> Optional["Tag"]: if value_offset + length > current_size: # we coerce the chunk size to be at least the size of the tag - chunk_size = max(value_offset + length - current_size, HEADER_CHUNK_SIZE) + chunk_size = max(value_offset + length - current_size, config.HEADER_CHUNK_SIZE) reader.data += await reader.range_request(len(reader.data), chunk_size, is_header=True) # read the tag value diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 806077f..082e28a 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -3,6 +3,7 @@ import pytest import rasterio +from aiocogeo import config from aiocogeo.ifd import IFD from aiocogeo.tag import BaseTag from aiocogeo.errors import InvalidTiffError @@ -168,4 +169,20 @@ async def test_cog_not_a_tiff(create_cog_reader): async def test_file_not_found(create_cog_reader, infile): with pytest.raises(FileNotFoundError): async with create_cog_reader(infile) as cog: - ... \ No newline at end of file + ... + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunk_size,request_count,header_size", [ + [16384, 2, 32770], + [4096, 4, 28769], + [100, 14, 26776] + ] +) +async def test_chunk_size(chunk_size, request_count, header_size, monkeypatch, create_cog_reader): + monkeypatch.setattr(config, "HEADER_CHUNK_SIZE", chunk_size) + async with create_cog_reader("https://async-cog-reader-test-data.s3.amazonaws.com/webp_cog.tif") as cog: + requests = cog.requests + assert requests['count'] == request_count + assert requests['header_size'] == header_size \ No newline at end of file