import unittest

import mock

import chainer
from chainer import testing
from chainer.testing import attr
from chainer.functions.pooling import pooling_nd_kernel


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

    def setUp(self):
        chainer.cuda.clear_memo()

    def test_pooling_nd_kernel_forward_memo(self):
        ndim = self.ndim
        with mock.patch('chainer.functions.pooling.pooling_nd_kernel.'
                        'PoolingNDKernelForward._generate') as m:
            pooling_nd_kernel.PoolingNDKernelForward.generate(ndim)
            m.assert_called_once_with(ndim)
            pooling_nd_kernel.PoolingNDKernelForward.generate(ndim)
            m.assert_called_once_with(ndim)

    def test_pooling_nd_kernel_backward_memo(self):
        ndim = self.ndim
        with mock.patch('chainer.functions.pooling.pooling_nd_kernel.'
                        'PoolingNDKernelBackward._generate') as m:
            pooling_nd_kernel.PoolingNDKernelBackward.generate(ndim)
            m.assert_called_once_with(ndim)
            pooling_nd_kernel.PoolingNDKernelBackward.generate(ndim)
            m.assert_called_once_with(ndim)


testing.run_module(__name__, __file__)

Add test for memoization in N-dimensional pooling kernel generator.