Update the TestDiskCachedDecorator test for collective

test/unit/test_caching.py

@@ -39,6 +39,7 @@
 import numpy
 from pyop2 import op2, mpi
 from pyop2.caching import disk_cached
+from pyop2.mpi import comm_cache_keyval
 
 
 def _seed():
@@ -555,20 +556,23 @@ def test_decorator_in_memory_cache_reuses_results(self, cache, cachedir):
         assert len(cache) == 1
         assert len(os.listdir(cachedir.name)) == 1
 
-    def test_decorator_collective_has_different_in_memory_key(self, cache, cachedir):
+    def test_decorator_collective_uses_different_in_memory_caches(self, cache, cachedir):
         decorated_func = disk_cached(cache, cachedir.name)(self.myfunc)
         collective_func = disk_cached(cache, cachedir.name, self.collective_key,
                                       collective=True)(self.myfunc)
 
+        # obj1 should be cached on the comm cache and not the self.cache
         obj1 = collective_func("input1")
-        assert len(cache) == 1
+        comm_cache = self.comm.Get_attr(comm_cache_keyval)
+        assert len(cache) == 0
+        assert len(comm_cache) == 1
         assert len(os.listdir(cachedir.name)) == 1
 
-        # The new entry should have a different in-memory key since the communicator
-        # is not included but the same key on disk.
+        # obj2 should be cached on the self.cache and not the comm cache
         obj2 = decorated_func("input1")
         assert obj1 == obj2 and obj1 is not obj2
-        assert len(cache) == 2
+        assert len(cache) == 1
+        assert len(comm_cache) == 1
         assert len(os.listdir(cachedir.name)) == 1
 
     def test_decorator_disk_cache_reuses_results(self, cache, cachedir):