import numpy as np
from numba import njit
from numba import unittest_support as unittest
from numba.tests import support
from numba.errors import TypingError

import mymodule
import myjitmodule2 # noqa - has side-effect, overload mymodule.set_to_x


@njit
def wrap_set_to_x(arr, x):
    mymodule.set_to_x(arr, x)


class TestSetToX(support.TestCase):

    def test_int(self):
        a = np.arange(10)
        wrap_set_to_x(a, 1)
        self.assertPreciseEqual(np.ones(10, dtype=np.int64), a)

    def test_float(self):
        a = np.arange(10, dtype=np.float64)
        wrap_set_to_x(a, 1.0)
        self.assertPreciseEqual(np.ones(10), a)

    def test_float_exception_on_nan(self):
        a = np.arange(10, dtype=np.float64)
        a[0] = np.nan
        with self.assertRaises(ValueError) as e:
            wrap_set_to_x(a, 1.0)
        self.assertIn("no element of arr may be NaN",
                      str(e.exception))

    def test_type_mismatch(self):
        a = np.arange(10)
        with self.assertRaises(TypingError) as e:
            wrap_set_to_x(a, 1.0)
        self.assertIn("the types of the inputs do not match",
                      str(e.exception))

    def test_exception_on_unsupported_dtype(self):
        a = np.arange(10, dtype=np.complex128)
        with self.assertRaises(TypingError) as e:
            wrap_set_to_x(a, np.complex128(1.0))
        self.assertIn("only integer and floating-point types are allowed",
                      str(e.exception))

    def test_exception_on_tuple(self):
        a = (1, 2, 3)
        with self.assertRaises(TypingError) as e:
            wrap_set_to_x(a, 1)
        self.assertIn("tuple isn't allowed as input, use NumPy ndarray",
                      str(e.exception))


if __name__ == '__main__':
    unittest.main()

@@ -357,48 +357,57 @@
 def 
-test_int(self):%0A        a = np.arange(10
+_check(self, a, b, x):%0A        wrap_set_to_x(a, x
 )%0A  
@@ -417,37 +417,45 @@
    wrap_set_to_x
-(a, 1
+.py_func(b, x
 )%0A        self.a
@@ -476,38 +476,160 @@
 ual(
-np.ones(10, dtype=np.int64), a
+a, b)%0A%0A    def test_int(self):%0A        self._check(np.zeros(10, dtype=np.int64),%0A                    np.zeros(10, dtype=np.int64),%0A                    1
 )%0A%0A 
@@ -657,37 +657,44 @@
 f):%0A        
-a = np.arange
+self._check(np.zeros
 (10, dtype=n
@@ -695,32 +695,33 @@
 type=np.float64)
+,
 %0A        wrap_se
@@ -717,76 +717,75 @@
     
-wrap_set_to_x(a, 1.0)%0A        self.assertPreciseEqual(np.ones(10), a
+            np.zeros(10, dtype=np.float64),%0A                    1.0
 )%0A%0A 

fix tests: check against orginal