tests: Tighten sparse-op Call assertions to exact names

Author: mloubout

tests/test_dle.py

@@ -180,9 +180,8 @@ def test_cache_blocking_structure_distributed(mode):
     # The sparse inject lives in its own efunc, which sits between the
     # two dense Eq's. Because they can no longer fuse across the
     # sparse-op Call, each dense Eq lands in its own MPI compute efunc.
-    compute_names = sorted(n for n in op._func_table if n.startswith('compute'))
-    bns0, _ = assert_blocking(op._func_table[compute_names[0]].root, {'x0_blk0'})
-    bns1, _ = assert_blocking(op._func_table[compute_names[1]].root, {'x1_blk0'})
+    bns0, _ = assert_blocking(op._func_table['compute0'].root, {'x0_blk0'})
+    bns1, _ = assert_blocking(op._func_table['compute1'].root, {'x1_blk0'})
 
     for blk_dim, bns in [('x0_blk0', bns0), ('x1_blk0', bns1)]:
         iters = FindNodes(Iteration).visit(bns[blk_dim])

tests/test_interpolation.py

@@ -995,8 +995,7 @@ def test_interp_complex_and_real(self, dtype):
         # Both interpolations land in the same sparse-op efunc since they
         # share the `p_sc` Dimension (sce reuses sc's coordinates); two
         # radius nests sit side-by-side inside the single ``p_sc`` loop.
-        [efunc_name] = [n for n in opC._func_table if n.startswith('interpolate_')]
-        efunc = opC._func_table[efunc_name].root
+        efunc = opC._func_table['interpolate_sc0'].root
         assert_structure(
             efunc,
             ['p_sc', 'p_sc,rp_scx,rp_scy,rp_scz', 'p_sc,rp_scx,rp_scy,rp_scz'],

tests/test_mpi.py

@@ -804,11 +804,10 @@ class SparseFirst(SparseFunction):
         rec = s.interpolate(expr=s+fs, implicit_dims=grid.stepping_dim)
         op = Operator(eqs + rec)
 
-        # The sparse interp now lowers to a Call into an
-        # ElementalFunction, so the parent carries that Call in
-        # addition to any halo exchanges.
-        calls = FindNodes(Call).visit(op)
-        assert any(c.name.startswith('interpolate_') for c in calls)
+        # Generated code: one halo exchange for ``fs`` and one Call
+        # to the ``interpolate_s0`` ElementalFunction.
+        call_names = [c.name for c in FindNodes(Call).visit(op)]
+        assert call_names == ['haloupdate0', 'interpolate_s0']
 
         op(time_M=10)
         expected = 10*11/2  # n (n+1)/2
@@ -838,10 +837,10 @@ class CoordSlowSparseFunction(SparseFunction):
 
         op = Operator([Eq(u, 1)] + rec_eq)
 
-        # Expected Calls: one halo exchange + one interpolate efunc.
-        call_names = sorted(c.name for c in FindNodes(Call).visit(op))
-        assert any(n.startswith('haloupdate') for n in call_names)
-        assert any(n.startswith('interpolate_') for n in call_names)
+        # Generated code: one halo exchange for ``u`` and one Call to
+        # the ``interpolate_s0`` ElementalFunction.
+        call_names = [c.name for c in FindNodes(Call).visit(op)]
+        assert call_names == ['haloupdate0', 'interpolate_s0']
 
         op.apply()
         assert np.all(s.data == 1)
@@ -870,10 +869,10 @@ class CoordSlowSparseFunction(SparseTimeFunction):
 
         op = Operator([Eq(u, 1)] + rec_eq)
 
-        # Expected Calls: one halo exchange + one interpolate efunc.
-        call_names = sorted(c.name for c in FindNodes(Call).visit(op))
-        assert any(n.startswith('haloupdate') for n in call_names)
-        assert any(n.startswith('interpolate_') for n in call_names)
+        # Generated code: one halo exchange for ``u`` and one Call to
+        # the ``interpolate_s0`` ElementalFunction.
+        call_names = [c.name for c in FindNodes(Call).visit(op)]
+        assert call_names == ['haloupdate0', 'interpolate_s0']
 
         op.apply(time_M=5)
         assert np.all(s.data == 1)