From 05de1b5835737601b37cae1c82f23025700efc1d Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Thu, 21 Jun 2018 19:26:39 -0500
Subject: [PATCH] Adapt test_preamble_with_separate_temporaries to make sure
 registered callables unpickle equally

---
 test/test_loopy.py | 12 ++++++++----
 test/testlib.py    | 26 +++++++++++++++++++++++---
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/test/test_loopy.py b/test/test_loopy.py
index 07179ae14..3ceca5a75 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -2719,9 +2719,13 @@ def test_preamble_with_separate_temporaries(ctx_factory):
         read_only=True),
      lp.GlobalArg('data', shape=(data.size,), dtype=np.float64)],
     )
+
     # fixt params, and add manglers / preamble
-    from testlib import SeparateTemporariesPreambleTestHelper
-    preamble_with_sep_helper = SeparateTemporariesPreambleTestHelper(
+    from testlib import (
+            SeparateTemporariesPreambleTestMangler,
+            SeparateTemporariesPreambleTestPreambleGenerator,
+            )
+    func_info = dict(
             func_name='indirect',
             func_arg_dtypes=(np.int32, np.int32, np.int32),
             func_result_dtypes=(np.int32,),
@@ -2730,9 +2734,9 @@ def test_preamble_with_separate_temporaries(ctx_factory):
 
     kernel = lp.fix_parameters(kernel, **{'n': n})
     kernel = lp.register_preamble_generators(
-            kernel, [preamble_with_sep_helper.preamble_gen])
+            kernel, [SeparateTemporariesPreambleTestPreambleGenerator(**func_info)])
     kernel = lp.register_function_manglers(
-            kernel, [preamble_with_sep_helper.mangler])
+            kernel, [SeparateTemporariesPreambleTestMangler(**func_info)])
 
     print(lp.generate_code(kernel)[0])
     # and call (functionality unimportant, more that it compiles)
diff --git a/test/testlib.py b/test/testlib.py
index 73de4199d..399b4852a 100644
--- a/test/testlib.py
+++ b/test/testlib.py
@@ -1,5 +1,7 @@
 import loopy as lp
 
+from collections import namedtuple
+
 
 # {{{ test_barrier_in_overridden_get_grid_size_expanded_kernel
 
@@ -17,14 +19,29 @@ class GridOverride(object):
 
 # {{{ test_preamble_with_separate_temporaries
 
-class SeparateTemporariesPreambleTestHelper:
+class SeparateTemporariesPreambleTestDataHolder:
     def __init__(self, func_name, func_arg_dtypes, func_result_dtypes, arr):
         self.func_name = func_name
         self.func_arg_dtypes = func_arg_dtypes
         self.func_result_dtypes = func_result_dtypes
         self.arr = arr
 
-    def mangler(self, kernel, name, arg_dtypes):
+    def __eq__(self, other):
+        import numpy as np
+        return (
+                isinstance(other, type(self))
+                and self.func_name == other.func_name
+                and self.func_arg_dtypes == other.func_arg_dtypes
+                and self.func_result_dtypes == other.func_result_dtypes
+                and np.array_equal(self.arr, other.arr))
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+class SeparateTemporariesPreambleTestMangler(
+        SeparateTemporariesPreambleTestDataHolder):
+    def __call__(self, kernel, name, arg_dtypes):
         """
         A function that will return a :class:`loopy.kernel.data.CallMangleInfo`
         to interface with the calling :class:`loopy.LoopKernel`
@@ -61,7 +78,10 @@ class SeparateTemporariesPreambleTestHelper:
                                 self.func_result_dtypes),
             arg_dtypes=arg_dtypes)
 
-    def preamble_gen(self, preamble_info):
+
+class SeparateTemporariesPreambleTestPreambleGenerator(
+        SeparateTemporariesPreambleTestDataHolder):
+    def __call__(self, preamble_info):
         from loopy.kernel.data import temp_var_scope as scopes
 
         # find a function matching our name
-- 
GitLab