Skip to content

Commit

Permalink
Add fixture to cleanup dataloader after each test runs (#1852)
Browse files Browse the repository at this point in the history
  • Loading branch information
oliverholworthy committed Jul 4, 2023
1 parent aad1112 commit aa5c25c
Showing 1 changed file with 17 additions and 0 deletions.
17 changes: 17 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import subprocess
import time
from pathlib import Path
from unittest.mock import patch

import dask
import pandas as pd
Expand Down Expand Up @@ -371,3 +372,19 @@ def devices(request):
@pytest.fixture
def report(request):
return request.config.getoption("--report")


@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).
"""
from merlin.dataloader.loader_base import LoaderBase

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()

0 comments on commit aa5c25c

Please sign in to comment.