import unittest

import mock

import chainer
from chainer import testing
from chainer.testing import attr
from chainer.utils import conv_nd_kernel


@testing.parameterize(*testing.product({
    'ndim': [2, 3, 4],
}))
@attr.gpu
class TestIm2colNDKernelMemo(unittest.TestCase):

    def test_im2col_nd_kernel_memo(self):
        ndim = self.ndim
        with mock.patch(
                'chainer.utils.conv_nd_kernel.Im2colNDKernel._generate') as m:
            chainer.cuda.clear_memo()
            conv_nd_kernel.Im2colNDKernel.generate(ndim)
            m.assert_called_once_with(ndim)
            conv_nd_kernel.Im2colNDKernel.generate(ndim)
            m.assert_called_once_with(ndim)


@testing.parameterize(*testing.product({
    'ndim': [2, 3, 4],
}))
@attr.gpu
class TestCol2imNDKernelMemo(unittest.TestCase):

    def test_col2im_nd_kernel_memo(self):
        ndim = self.ndim
        with mock.patch(
                'chainer.utils.conv_nd_kernel.Col2imNDKernel._generate') as m:
            chainer.cuda.clear_memo()
            conv_nd_kernel.Col2imNDKernel.generate(ndim)
            m.assert_called_once_with(ndim)
            conv_nd_kernel.Col2imNDKernel.generate(ndim)
            m.assert_called_once_with(ndim)


testing.run_module(__name__, __file__)

@@ -262,32 +262,88 @@
 est.TestCase):%0A%0A
+    def setUp(self):%0A        chainer.cuda.clear_memo()%0A%0A
     def test_im2
@@ -360,32 +360,32 @@
 nel_memo(self):%0A
-
         ndim = s
@@ -501,46 +501,8 @@
  m:%0A
-            chainer.cuda.clear_memo()%0A
     
@@ -821,32 +821,88 @@
 est.TestCase):%0A%0A
+    def setUp(self):%0A        chainer.cuda.clear_memo()%0A%0A
     def test_col
@@ -969,32 +969,32 @@
 ith mock.patch(%0A
+
                 
@@ -1060,46 +1060,8 @@
  m:%0A
-            chainer.cuda.clear_memo()%0A
     

Move clearing memo to setup method.