python/utils/dataset_test.py

0001 from builtins import range
0002 from .dataset import *
0003
0004 import unittest
0005 import os
0006 import shutil
0007
0008 BASE_DIR = 'datasets'
0009
0010 def create_dataset(name, number_of_files, basedir=BASE_DIR):
0011     if not os.path.isdir(basedir):
0012         os.mkdir(basedir)
0013     old = os.getcwd()
0014     os.chdir(basedir)
0015     if os.path.isdir(name):
0016         shutil.rmtree(name)
0017     os.mkdir(name)
0018     os.chdir(name)
0019     for i in range(number_of_files):
0020         os.system('touch file_{i:d}.root'.format(i=i))
0021     os.chdir(old)
0022
0023 class TestDataset(unittest.TestCase):
0024
0025     def test_local(self):
0026         n_files = 10
0027         create_dataset('ds1', n_files)
0028         ds1 = LocalDataset('/ds1', 'datasets', '.*root')
0029         self.assertEqual( len(ds1.listOfGoodFiles()), n_files)
0030         # shutil.rmtree('datasets')
0031
0032     def test_eos(self):
0033         ds1 = EOSDataset('/EOSTests/ds1',
0034                          '/eos/cms/store/cmst3/user/cbern',
0035                          '.*root')
0036         self.assertEqual(len(ds1.listOfGoodFiles()), 10)
0037
0038     def test_eos_fail(self):
0039         self.assertRaises( ValueError,
0040                            EOSDataset,
0041                            'not_existing_basedir',
0042                            'not_exiting_dataset',
0043                            '.*root')
0044         # should test that we fail when a plain file is provided
0045         # instead of a directory.. but eostools not set up for that yet.
0046
0047
0048 if __name__ == '__main__':
0049     import os
0050     import sys
0051     if not os.environ.get('CMSSW_BASE', False):
0052         sys.exit(1)
0053     unittest.main()