python/utils/dataset_test.py

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