def test_remove_add_pass(self):
"""Test ImpactFunc can be added after removing."""
imp_fun = ImpactFuncSet()
vulner_1 = ImpactFunc()
vulner_1.id = 1
vulner_1.haz_type = 'TC'
imp_fun.append(vulner_1)
imp_fun.remove_func()
self.assertEqual(0, len(imp_fun.get_hazard_types()))
self.assertEqual(0, len(imp_fun.get_ids()))
imp_fun.append(vulner_1)
self.assertEqual(1, len(imp_fun.get_hazard_types()))
self.assertEqual('TC', imp_fun.get_hazard_types()[0])
self.assertEqual(1, len(imp_fun.get_ids()))
self.assertEqual([1], imp_fun.get_ids('TC'))
def test_get_ids_wrong_zero(self):
"""Test get_ids method with wrong inputs."""
imp_fun = ImpactFuncSet()
vulner_1 = ImpactFunc()
vulner_1.haz_type = 'WS'
vulner_1.id = 56
imp_fun.append(vulner_1)
self.assertEqual([], imp_fun.get_ids('TC'))
def test_get_ids_pass(self):
"""Test normal functionality of get_ids method."""
imp_fun = ImpactFuncSet()
self.assertEqual({}, imp_fun.get_ids())
vulner_1 = ImpactFunc()
vulner_1.id = 1
vulner_1.haz_type = 'TC'
imp_fun.append(vulner_1)
self.assertEqual(1, len(imp_fun.get_ids()))
self.assertIn('TC', imp_fun.get_ids())
self.assertEqual(1, len(imp_fun.get_ids('TC')))
self.assertEqual([1], imp_fun.get_ids('TC'))
vulner_2 = ImpactFunc()
vulner_2.id = 3
vulner_2.haz_type = 'TC'
imp_fun.append(vulner_2)
self.assertEqual(1, len(imp_fun.get_ids()))
self.assertIn('TC', imp_fun.get_ids())
self.assertEqual(2, len(imp_fun.get_ids('TC')))
self.assertEqual([1, 3], imp_fun.get_ids('TC'))
vulner_3 = ImpactFunc()
vulner_3.id = 3
vulner_3.haz_type = 'FL'
imp_fun.append(vulner_3)
self.assertEqual(2, len(imp_fun.get_ids()))
self.assertIn('TC', imp_fun.get_ids())
self.assertIn('FL', imp_fun.get_ids())
self.assertEqual(2, len(imp_fun.get_ids('TC')))
self.assertEqual([1, 3], imp_fun.get_ids('TC'))
self.assertEqual(1, len(imp_fun.get_ids('FL')))
self.assertEqual([3], imp_fun.get_ids('FL'))
