From 78df9bb5cf13b3a8baac920e8e405949ba20a5c8 Mon Sep 17 00:00:00 2001 From: Joan Fontanals Martinez Date: Thu, 28 Sep 2023 18:46:26 +0200 Subject: [PATCH] test: move the pydantic check inside test Signed-off-by: Joan Fontanals Martinez --- docarray/store/helpers.py | 54 ++++++++++++++------------- tests/integrations/store/test_file.py | 27 ++++++-------- tests/integrations/store/test_s3.py | 2 +- 3 files changed, 42 insertions(+), 41 deletions(-) diff --git a/docarray/store/helpers.py b/docarray/store/helpers.py index 35c6fd2c6ae..e2c4cf99a5d 100644 --- a/docarray/store/helpers.py +++ b/docarray/store/helpers.py @@ -174,32 +174,36 @@ def _from_binary_stream( compress: Optional[str] = None, show_progress: bool = False, ) -> Iterator['T']: - if show_progress: - pbar, t = _get_progressbar( - 'Deserializing', disable=not show_progress, total=total - ) - else: - pbar = nullcontext() - - with pbar: + try: if show_progress: - _total_size = 0 - pbar.start_task(t) - while True: - len_bytes = stream.read(4) - if len(len_bytes) < 4: - raise ValueError('Unexpected end of stream') - len_item = int.from_bytes(len_bytes, 'big', signed=False) - if len_item == 0: - break - item_bytes = stream.read(len_item) - if len(item_bytes) < len_item: - raise ValueError('Unexpected end of stream') - item = cls.from_bytes(item_bytes, protocol=protocol, compress=compress) - - yield item + pbar, t = _get_progressbar( + 'Deserializing', disable=not show_progress, total=total + ) + else: + pbar = nullcontext() + with pbar: if show_progress: - _total_size += len_item + 4 - pbar.update(t, advance=1, total_size=str(filesize.decimal(_total_size))) + _total_size = 0 + pbar.start_task(t) + while True: + len_bytes = stream.read(4) + if len(len_bytes) < 4: + raise ValueError('Unexpected end of stream') + len_item = int.from_bytes(len_bytes, 'big', signed=False) + if len_item == 0: + break + item_bytes = stream.read(len_item) + if len(item_bytes) < len_item: + raise ValueError('Unexpected end of stream') + item = cls.from_bytes(item_bytes, protocol=protocol, compress=compress) + + yield item + + if show_progress: + _total_size += len_item + 4 + pbar.update( + t, advance=1, total_size=str(filesize.decimal(_total_size)) + ) + finally: stream.close() diff --git a/tests/integrations/store/test_file.py b/tests/integrations/store/test_file.py index 87c7b2ee3f2..c222f6bdc61 100644 --- a/tests/integrations/store/test_file.py +++ b/tests/integrations/store/test_file.py @@ -6,12 +6,12 @@ from docarray import DocList from docarray.documents import TextDoc from docarray.store.file import ConcurrentPushException, FileDocStore -from docarray.utils._internal.cache import _get_cache_path from docarray.utils._internal.pydantic import is_pydantic_v2 +from docarray.utils._internal.cache import _get_cache_path from tests.integrations.store import gen_text_docs, get_test_da, profile_memory DA_LEN: int = 2**10 -TOLERANCE_RATIO = 0.1 # Percentage of difference allowed in stream vs non-stream test +TOLERANCE_RATIO = 0.1 # Percentage of difference allowed when streaming between a long and a shorter DA def test_path_resolution(): @@ -23,7 +23,6 @@ def test_path_resolution(): def test_pushpull_correct(capsys, tmp_path: Path): - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = tmp_path da1 = get_test_da(DA_LEN) @@ -51,7 +50,6 @@ def test_pushpull_correct(capsys, tmp_path: Path): def test_pushpull_stream_correct(capsys, tmp_path: Path): - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = tmp_path da1 = get_test_da(DA_LEN) @@ -85,10 +83,8 @@ def test_pushpull_stream_correct(capsys, tmp_path: Path): # for some reason this test is failing with pydantic v2 -@pytest.mark.skipif(is_pydantic_v2, reason="Not working with pydantic v2 for now") @pytest.mark.slow def test_pull_stream_vs_pull_full(tmp_path: Path): - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = tmp_path DocList[TextDoc].push_stream( gen_text_docs(DA_LEN * 1), @@ -135,16 +131,19 @@ def get_total_full(url: str): long_total_stream == long_total_full ), 'Streamed and non-streamed pull should have similar statistics' - assert ( - abs(long_stream_peak - short_stream_peak) / short_stream_peak < TOLERANCE_RATIO - ), 'Streamed memory usage should not be dependent on the size of the data' - assert ( - abs(long_full_peak - short_full_peak) / short_full_peak > TOLERANCE_RATIO - ), 'Full pull memory usage should be dependent on the size of the data' + assert long_full_peak > long_stream_peak, 'Peak of memory using full should be larger than when streaming' + assert short_full_peak > short_stream_peak, 'Peak of memory using full should be larger than when streaming' + if not is_pydantic_v2: + # I bet there is some memory that Pydantic is leaking + assert ( + abs(long_stream_peak - short_stream_peak) / short_stream_peak < TOLERANCE_RATIO + ), 'Streamed memory usage should not be dependent on the size of the data' + assert ( + abs(long_full_peak - short_full_peak) / short_full_peak > TOLERANCE_RATIO + ), 'Full pull memory usage should be dependent on the size of the data' def test_list_and_delete(tmp_path: Path): - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = str(tmp_path) da_names = FileDocStore.list(namespace_dir, show_table=False) @@ -179,7 +178,6 @@ def test_list_and_delete(tmp_path: Path): def test_concurrent_push_pull(tmp_path: Path): # Push to DA that is being pulled should not mess up the pull - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = tmp_path DocList[TextDoc].push_stream( @@ -214,7 +212,6 @@ def test_concurrent_push(tmp_path: Path): # Double push should fail the second push import time - tmp_path.mkdir(parents=True, exist_ok=True) namespace_dir = tmp_path DocList[TextDoc].push_stream( diff --git a/tests/integrations/store/test_s3.py b/tests/integrations/store/test_s3.py index eb36ddd907a..a7c8d4c5e9e 100644 --- a/tests/integrations/store/test_s3.py +++ b/tests/integrations/store/test_s3.py @@ -12,7 +12,7 @@ from tests.integrations.store import gen_text_docs, get_test_da, profile_memory DA_LEN: int = 2**10 -TOLERANCE_RATIO = 0.5 # Percentage of difference allowed in stream vs non-stream test +TOLERANCE_RATIO = 0.1 # Percentage of difference allowed in stream vs non-stream test BUCKET: str = 'da-pushpull' RANDOM: str = uuid.uuid4().hex[:8]