From fdd3bf57125b7a090bed5cd15e58fca2913140eb Mon Sep 17 00:00:00 2001
From: Davor Dundovic <33790330+ddundo@users.noreply.github.com>
Date: Thu, 24 Oct 2024 15:51:23 +0200
Subject: [PATCH] #222: Update caching (#225)

Closes #222.
---
 test/conftest.py         | 9 +++++----
 test_adjoint/conftest.py | 9 +++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/test/conftest.py b/test/conftest.py
index bb45c5cf..3675a4e0 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -21,8 +21,9 @@ def pytest_runtest_teardown(item, nextitem):
     """
     Clear caches after running a test
     """
-    from firedrake.tsfc_interface import TSFCKernel
-    from pyop2.global_kernel import GlobalKernel
+    from firedrake.tsfc_interface import clear_cache
+    from pyop2.caching import clear_memory_cache
+    from pyop2.mpi import COMM_WORLD
 
-    TSFCKernel._cache.clear()
-    GlobalKernel._cache.clear()
+    clear_cache()
+    clear_memory_cache(COMM_WORLD)
diff --git a/test_adjoint/conftest.py b/test_adjoint/conftest.py
index 8a582c51..daf0543e 100644
--- a/test_adjoint/conftest.py
+++ b/test_adjoint/conftest.py
@@ -135,11 +135,12 @@ def pytest_runtest_teardown(item, nextitem):
     """
     Clear caches after running a test
     """
-    from firedrake.tsfc_interface import TSFCKernel
-    from pyop2.global_kernel import GlobalKernel
+    from firedrake.tsfc_interface import clear_cache
+    from pyop2.caching import clear_memory_cache
+    from pyop2.mpi import COMM_WORLD
 
-    TSFCKernel._cache.clear()
-    GlobalKernel._cache.clear()
+    clear_cache()
+    clear_memory_cache(COMM_WORLD)
 
 
 @pytest.fixture(autouse=True)