Paddle/python/paddle/v2/dataset/tests/mnist_test.py

import paddle.v2.dataset.mnist
import unittest

class TestMNIST(unittest.TestCase):
    def check_reader(self, reader):
        sum = 0
        for l in reader:
            self.assertEqual(l[0].size, 784)
            self.assertEqual(l[1].size, 1)
            self.assertLess(l[1], 10)
            self.assertGreaterEqual(l[1], 0)
            sum += 1
        return sum

    def test_train(self):
        self.assertEqual(
            self.check_reader(paddle.v2.dataset.mnist.train()),
            60000)

    def test_test(self):
        self.assertEqual(
            self.check_reader(paddle.v2.dataset.mnist.test()),
            10000)


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