Skip to content

Commit d1a14de

Browse files
author
Victor Machado
committed
OHSU-32 More froster config tests added. Pending slurm testing
1 parent 44e834d commit d1a14de

File tree

2 files changed

+125
-4
lines changed

2 files changed

+125
-4
lines changed

Diff for: ‎froster/froster.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -1004,6 +1004,10 @@ def set_shared(self):
10041004
shared_config_dir, self.shared_config_file_name))
10051005
self.__set_shared_move_config()
10061006

1007+
self.is_shared = True
1008+
else:
1009+
self.is_shared = False
1010+
10071011
print(f'*** SHARED CONFIGURATION DONE ***\n')
10081012

10091013
return True
@@ -5612,7 +5616,9 @@ def __init__(self, args, cfg: ConfigManager):
56125616
self.squeue_output_format = '"%i","%j","%t","%M","%L","%D","%C","%m","%b","%R"'
56135617
self.jobs = []
56145618
self.job_info = {}
5619+
56155620
self.partition = cfg.slurm_partition
5621+
56165622
self.qos = cfg.slurm_qos
56175623
self.walltime = f'{cfg.slurm_walltime_days}-{cfg.slurm_walltime_hours}'
56185624
self._add_lines_from_cfg()
@@ -5778,7 +5784,7 @@ def _get_associations(self):
57785784
",") for item in self._parse_tabular_data(mystr) if 'Account' in item}
57795785
return asso
57805786

5781-
def get_allowed_partitions_and_qos(self, account=None):
5787+
def get_allowed_partitions_and_qos(self):
57825788
"""Get a dictionary with keys = partitions and values = QOSs the user is allowed to use."""
57835789
bacc = os.environ.get('SBATCH_ACCOUNT', '')
57845790
account = bacc if bacc else account

Diff for: ‎tests/test_config.py

+118-3
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,16 @@
3434
S3_STORAGE_CLASS = 'DEEP_ARCHIVE'
3535
S3_STORAGE_CLASS_2 = 'GLACIER'
3636

37+
SLURM_WALLTIME_DAYS = 8
38+
SLURM_WALLTIME_HOURS = 1
39+
SLURM_PARTITION = 'test_partition'
40+
SLURM_QOS = 'test_qos'
41+
42+
SLURM_LOCAL_SCRATCH = 'test_lscratch'
43+
SLURM_SCRIPT_SCRATCH = 'test_script_scratch'
44+
SLURM_SCRIPT_TEARS_DOWN = 'test_script_tears_down'
45+
SLURM_ROOT = 'test_root'
46+
3747
SHARED_DIR = os.path.join(tempfile.gettempdir(), 'shared_dir')
3848

3949

@@ -493,7 +503,7 @@ def test_set_shared_froster_shared_config_exist(self, mock_print, mock_input_con
493503
self.assertTrue(os.path.exists(archive_json_file_shared))
494504

495505

496-
@patch('builtins.print')
506+
# @patch('builtins.print')
497507
class TestConfigNIH(unittest.TestCase):
498508

499509
# Method executed before every test
@@ -532,6 +542,28 @@ def test_set_nih(self, mock_print, mock_input_confirm):
532542
NIH_SECTION, 'is_nih', 'False')
533543

534544

545+
@patch('inquirer.confirm', side_effect=[True, True, False])
546+
@patch('inquirer.prompt', return_value={'shared_dir': SHARED_DIR})
547+
def test_set_nih_when_shared_config(self, mock_print, mock_input_confirm, mock_input_promp):
548+
'''- Set the NIH flag to True and then to False.'''
549+
550+
# Call set_shared method
551+
self.assertTrue(self.cfg.set_shared())
552+
553+
# Call set_user method
554+
self.assertTrue(self.cfg.set_nih())
555+
556+
# Check that the configuration file was updated correctly
557+
check_ini_file(self, self.cfg.shared_config_file,
558+
NIH_SECTION, 'is_nih', 'True')
559+
560+
# Call set_user method
561+
self.assertTrue(self.cfg.set_nih())
562+
563+
# Check that the configuration file was updated correctly
564+
check_ini_file(self, self.cfg.shared_config_file,
565+
NIH_SECTION, 'is_nih', 'False')
566+
535567
@patch('builtins.print')
536568
class TestConfigS3(unittest.TestCase):
537569

@@ -638,6 +670,39 @@ def test_set_s3(self, mock_print, mock_input_list, mock_input_text):
638670
# Check the bucket was created
639671
self.assertIn(S3_BUCKET_NAME, s3_buckets)
640672

673+
674+
@patch('inquirer.list_input', side_effect=['+ Create new bucket', S3_STORAGE_CLASS])
675+
@patch('inquirer.text', side_effect=[S3_BUCKET_NAME, S3_ARCHIVE_DIR])
676+
@patch('inquirer.confirm', side_effect=[True])
677+
@patch('inquirer.prompt', return_value={'shared_dir': SHARED_DIR})
678+
def test_set_s3_when_shared_config(self, mock_print, mock_input_list, mock_input_text, mock_input_confirm, mock_input_prompt):
679+
'''- Set a new S3 bucket when shared configuration.'''
680+
681+
# Call set_shared method
682+
self.assertTrue(self.cfg.set_shared())
683+
684+
# Call set_s3 method
685+
self.assertTrue(self.cfg.set_s3(self.aws))
686+
687+
# Check that the configuration files were updated correctly
688+
check_ini_file(self, self.cfg.shared_config_file,
689+
S3_SECTION, 'bucket_name', S3_BUCKET_NAME)
690+
691+
check_ini_file(self, self.cfg.shared_config_file,
692+
S3_SECTION, 'archive_dir', S3_ARCHIVE_DIR)
693+
694+
check_ini_file(self, self.cfg.shared_config_file,
695+
S3_SECTION, 'storage_class', S3_STORAGE_CLASS)
696+
697+
# Check that the s3_init is set
698+
self.assertTrue(self.cfg.s3_init)
699+
700+
# Get the buckets list
701+
s3_buckets = self.aws.get_buckets()
702+
703+
# Check the bucket was created
704+
self.assertIn(S3_BUCKET_NAME, s3_buckets)
705+
641706
@patch('inquirer.list_input', side_effect=['+ Create new bucket', S3_STORAGE_CLASS, S3_BUCKET_NAME, S3_STORAGE_CLASS_2])
642707
@patch('inquirer.text', side_effect=[S3_BUCKET_NAME, S3_ARCHIVE_DIR, S3_ARCHIVE_DIR_2])
643708
def test_set_s3_select_bucket(self, mock_print, mock_input_list, mock_input_text):
@@ -669,9 +734,57 @@ def test_set_s3_select_bucket(self, mock_print, mock_input_list, mock_input_text
669734
self.assertNotIn(S3_BUCKET_NAME, s3_buckets)
670735

671736

737+
# TODO: SLURM part pending
738+
739+
# # @patch('builtins.print')
740+
# class TestConfigSlurm(unittest.TestCase):
741+
742+
# # Method executed before every test
743+
# def setUp(self):
744+
# init_froster(self)
745+
746+
# # Method executed after every test
747+
# def tearDown(self):
748+
# deinit_froster(self)
749+
750+
# @patch('shutil.which', return_value=True)
751+
# @patch('inquirer.text', side_effect=[SLURM_WALLTIME_DAYS, SLURM_WALLTIME_HOURS, SLURM_LOCAL_SCRATCH, SLURM_SCRIPT_SCRATCH, SLURM_SCRIPT_TEARS_DOWN, SLURM_ROOT])
752+
# @patch('inquirer.list_input', side_effect=[SLURM_PARTITION, SLURM_QOS])
753+
# def test_set_slurm(self, mock_shutil, mock_input_text, mock_input_list):
754+
# '''- set_slurm'''
755+
756+
# # Call set_slurm method
757+
# self.cfg.set_slurm(self.args)
758+
759+
# # Check that the configuration files were updated correctly
760+
# check_ini_file(self, self.cfg.config_file,
761+
# 'SLURM', 'slurm_walltime_days', str(SLURM_WALLTIME_DAYS))
762+
763+
# check_ini_file(self, self.cfg.config_file,
764+
# 'SLURM', 'slurm_walltime_hours', str(SLURM_WALLTIME_HOURS))
765+
766+
# check_ini_file(self, self.cfg.config_file,
767+
# 'SLURM', 'slurm_partition', SLURM_PARTITION)
768+
769+
# check_ini_file(self, self.cfg.config_file,
770+
# 'SLURM', 'slurm_qos', SLURM_QOS)
771+
772+
# check_ini_file(self, self.cfg.config_file,
773+
# 'SLURM', 'slurm_lscratch', SLURM_LOCAL_SCRATCH)
774+
775+
# check_ini_file(self, self.cfg.config_file,
776+
# 'SLURM', 'lscratch_mkdir', SLURM_SCRIPT_SCRATCH)
777+
778+
# check_ini_file(self, self.cfg.config_file,
779+
# 'SLURM', 'lscratch_rmdir', SLURM_SCRIPT_TEARS_DOWN)
780+
781+
# check_ini_file(self, self.cfg.config_file,
782+
# 'SLURM', 'lscratch_root', SLURM_ROOT)
783+
784+
672785
if __name__ == '__main__':
673786

674-
if True:
787+
if False:
675788
unittest.main(verbosity=2)
676789
else:
677790
suite = unittest.TestSuite()
@@ -680,6 +793,8 @@ def test_set_s3_select_bucket(self, mock_print, mock_input_list, mock_input_text
680793
# suite.addTest(unittest.TestLoader().loadTestsFromTestCase(TestConfigShared))
681794
# suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestConfigNIH))
682795
# suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestConfigS3))
683-
# suite.addTest(TestConfigS3('test_set_s3_select_bucket'))
796+
# suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestConfigSlurm))
797+
# suite.addTest(TestConfigNIH('test_set_nih_when_shared_config'))
798+
suite.addTest(TestConfigS3('test_set_s3_when_shared_config'))
684799
runner = unittest.TextTestRunner(verbosity=2)
685800
runner.run(suite)

0 commit comments

Comments
 (0)