From b89f66e13dbce55dac4cf724f74b7a34cfe7b1c3 Mon Sep 17 00:00:00 2001 From: Oliver Holworthy Date: Sat, 1 Jul 2023 02:37:48 +0100 Subject: [PATCH] Add fixture to cleanup dataloader (#1167) Co-authored-by: edknv <109497216+edknv@users.noreply.github.com> --- tests/conftest.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 6e5daeb58b..386ea8884f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -20,6 +20,7 @@ import platform import warnings from pathlib import Path +from unittest.mock import patch import distributed import psutil @@ -27,6 +28,7 @@ from asvdb import BenchmarkInfo, utils from merlin.core.utils import Distributed +from merlin.dataloader.loader_base import LoaderBase from merlin.datasets.synthetic import generate_data from merlin.io import Dataset from merlin.models.utils import ci_utils @@ -145,3 +147,17 @@ def get_benchmark_info(): arch=uname.machine, ram="%d" % psutil.virtual_memory().total, ) + + +@pytest.fixture(scope="function", autouse=True) +def cleanup_dataloader(): + """After each test runs. Call .stop() on any dataloaders created during the test. + The avoids issues with background threads hanging around and interfering with subsequent tests. + This happens when a dataloader is partially consumed (not all batches are iterated through). + """ + with patch.object( + LoaderBase, "__iter__", side_effect=LoaderBase.__iter__, autospec=True + ) as patched: + yield + for call in patched.call_args_list: + call.args[0].stop()