Added tests.

STIXProject · Oct 21, 2016 · c5db4c9 · c5db4c9
1 parent 5126778
commit c5db4c9
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 16 deletions.
diff --git a/stix/extensions/marking/ais.py b/stix/extensions/marking/ais.py
@@ -443,18 +443,16 @@ def _update_schemalocations():
 
 def _validate_and_create_industry_type(industry_type):
     INDUSTRY_SECTORS = (CHEMICAL_SECTOR, COMMERCIAL_FACILITIES_SECTOR,
-                             COMMUNICATIONS_SECTOR,
-                             CRITICAL_MANUFACTURING_SECTOR,
-                             DAMS_SECTOR, DEFENSE_INDUSTRIAL_BASE_SECTOR,
-                             EMERGENCY_SERVICES_SECTOR, ENERGY_SECTOR,
-                             FINANCIAL_SERVICES_SECTOR,
-                             FOOD_AND_AGRICULTURE_SECTOR,
-                             GOVERNMENT_FACILITIES_SECTOR,
-                             HEALTH_CARE_AND_PUBLIC_HEALTH_SECTOR,
-                             INFORMATION_TECHNOLOGY_SECTOR,
-                             NUCLEAR_REACTORS_MATERIALS_AND_WASTE_SECTOR,
-                             TRANSPORTATION_SYSTEMS_SECTOR, OTHER,
-                             WATER_AND_WASTEWATER_SYSTEMS_SECTOR)
+                        COMMUNICATIONS_SECTOR, CRITICAL_MANUFACTURING_SECTOR,
+                        DAMS_SECTOR, DEFENSE_INDUSTRIAL_BASE_SECTOR,
+                        EMERGENCY_SERVICES_SECTOR, ENERGY_SECTOR,
+                        FINANCIAL_SERVICES_SECTOR, FOOD_AND_AGRICULTURE_SECTOR,
+                        GOVERNMENT_FACILITIES_SECTOR,
+                        HEALTH_CARE_AND_PUBLIC_HEALTH_SECTOR,
+                        INFORMATION_TECHNOLOGY_SECTOR,
+                        NUCLEAR_REACTORS_MATERIALS_AND_WASTE_SECTOR,
+                        TRANSPORTATION_SYSTEMS_SECTOR, OTHER,
+                        WATER_AND_WASTEWATER_SYSTEMS_SECTOR)
 
     lower_case_sectors = tuple(x.lower() for x in INDUSTRY_SECTORS)
     result = ""
@@ -485,7 +483,7 @@ def _validate_and_create_industry_type(industry_type):
             error = True
             break
 
-    if not error:
+    if not error and val:
         return result
 
     msg = 'IndustryType must be one of the following: {0}. Received \'{1}\'.'

diff --git a/stix/test/extensions/marking/ais.py b/stix/test/extensions/marking/ais.py
@@ -2,11 +2,12 @@
 
 from stix.test import EntityTestCase
 
-from stix.extensions.marking.ais import AISMarkingStructure
+from stix.core import STIXPackage
+from stix.extensions.marking import ais
 
 
 class AISMarkingStructureNotProprietaryTests(EntityTestCase, unittest.TestCase):
-    klass = AISMarkingStructure
+    klass = ais.AISMarkingStructure
     _full_dict = {
         'not_proprietary':
             {
@@ -19,7 +20,7 @@ class AISMarkingStructureNotProprietaryTests(EntityTestCase, unittest.TestCase):
 
 
 class AISMarkingStructureIsProprietaryTests(EntityTestCase, unittest.TestCase):
-    klass = AISMarkingStructure
+    klass = ais.AISMarkingStructure
     _full_dict = {
         'is_proprietary':
             {
@@ -30,5 +31,55 @@ class AISMarkingStructureIsProprietaryTests(EntityTestCase, unittest.TestCase):
         'xsi:type': 'AIS:AISMarkingStructure'
     }
 
+
+class AISMarkingGeneralTests(unittest.TestCase):
+    """General tests for ais module helpers."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.INDUSTRY_SECTORS = (ais.CHEMICAL_SECTOR, ais.COMMERCIAL_FACILITIES_SECTOR,
+                                ais.COMMUNICATIONS_SECTOR,
+                                ais.CRITICAL_MANUFACTURING_SECTOR,
+                                ais.DAMS_SECTOR, ais.DEFENSE_INDUSTRIAL_BASE_SECTOR,
+                                ais.EMERGENCY_SERVICES_SECTOR, ais.ENERGY_SECTOR,
+                                ais.FINANCIAL_SERVICES_SECTOR,
+                                ais.FOOD_AND_AGRICULTURE_SECTOR,
+                                ais.GOVERNMENT_FACILITIES_SECTOR,
+                                ais.HEALTH_CARE_AND_PUBLIC_HEALTH_SECTOR,
+                                ais.INFORMATION_TECHNOLOGY_SECTOR,
+                                ais.NUCLEAR_REACTORS_MATERIALS_AND_WASTE_SECTOR,
+                                ais.TRANSPORTATION_SYSTEMS_SECTOR, ais.OTHER,
+                                ais.WATER_AND_WASTEWATER_SYSTEMS_SECTOR)
+
+    def test_validate_and_create_industry_type(self):
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, [])
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, "")
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, [""])
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, ["Energy Sector", "Others", "Dams Sector"])
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, "Energy Sector| Others|Dams Sector")
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, 3)
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, "|")
+        self.assertRaises(ValueError, ais._validate_and_create_industry_type, ["Energy Sector|Dams Sector"])
+
+        self.assertEqual(ais._validate_and_create_industry_type("eNergY sectOr"), ais.ENERGY_SECTOR)
+        self.assertEqual(ais._validate_and_create_industry_type("eNergY sectOr   |Dams sector"), "Energy Sector|Dams Sector")
+        self.assertEqual(ais._validate_and_create_industry_type(["eNergY sectOr    ", "    Dams sectOr    "]), "Energy Sector|Dams Sector")
+        self.assertEqual(ais._validate_and_create_industry_type([ais.ENERGY_SECTOR, ais.DAMS_SECTOR]), "Energy Sector|Dams Sector")
+        self.assertEqual(ais._validate_and_create_industry_type("Energy Sector|Dams Sector"), "Energy Sector|Dams Sector")
+
+        for idx, x in enumerate(self.INDUSTRY_SECTORS):
+            self.assertEqual(ais._validate_and_create_industry_type(x), self.INDUSTRY_SECTORS[idx])
+
+    def test_add_ais_marking(self):
+        PACKAGE = STIXPackage()
+
+        # Missing kwarg.
+        self.assertRaises(ValueError, ais.add_ais_marking, PACKAGE, True, 'NONE', 'GREEN',
+                        country_name_code='US',
+                        country_name_code_type='ISO 3166-1 alpha-2',
+                        admin_area_name_code='US-DC',
+                        admin_area_name_code_type='ISO 3166-2',
+                        organisation_name='NCCIC')
+
 if __name__ == "__main__":
     unittest.main()