Optionally check NRT allocations (#286)

HISTORY.rst

@@ -4,6 +4,7 @@ History
 
 X.Y.Z (YYYY-MM-DD)
 ------------------
+* Optionally check NRT allocations (:pr:`286`)
 * Use `packaging.version.Version` instead of deprecated `distutils.version.LooseVersion` (:pr:`285`)
 * Pin numba to less than 0.59 in anticipation of @generated_jit deprecation (:pr:`284`)
 * Update trove hash (:pr:`279`)

africanus/averaging/tests/test_bda_averaging.py

@@ -368,23 +368,8 @@ def test_dask_bda_avg(vis_format):
             raise ValueError(f"Invalid vis_format: {vis_format}")
 
 
-@pytest.fixture
-def check_leaks():
-    import gc
-    from numba.core.runtime import rtsys
-
-    try:
-        yield None
-    finally:
-        gc.collect()
-
-    stats = rtsys.get_allocation_stats()
-    assert stats.alloc == stats.free
-    assert stats.mi_alloc == stats.mi_free
-
-
 @pytest.mark.parametrize("dtype", [np.complex64])
-def test_bda_output_arrays(dtype, check_leaks):
+def test_bda_output_arrays(dtype):
     from africanus.averaging.bda_avg import vis_output_arrays
     from numba import njit
 

africanus/conftest.py

@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
 
+import numba
 from numba.core.runtime import rtsys
 import pytest
 
 from africanus.util.testing import mark_in_pytest
 
 
-@pytest.fixture(scope="function", autouse=True)
+@pytest.fixture(scope="function", autouse=bool(numba.config.NRT_STATS))
 def check_allocations():
     """ Check allocations match frees """
     try: