

import unittest
import numpy as np
rng = np.random.RandomState()

from config_spec import DataSpec

class Test(unittest.TestCase):

    def testDataSpec(self):
        for i in range(5):
            a,b = DataSpec.batches_from_data(rng.randint(100, 1000), 5,
                                             rng.randint(1, 10) / 20., i, rng)
            self.assertEqual( list(set(a+b)), range(len(a+b)) )        
            self.assertTrue( i in b )



if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()