improve input validations

src/ua_generator/user_agent.py

@@ -14,9 +14,9 @@
 
 class UserAgent:
     def __init__(self, device=None, platform=None, browser=None, options=None):
-        self.device = device
-        self.platform = platform
-        self.browser = browser
+        self.device: str = utils.choice(device)
+        self.platform: str = utils.choice(platform)
+        self.browser: str = utils.choice(browser)
         self.options: Options = options if options is not None else Options()
         self.__complete()
 
@@ -26,17 +26,14 @@ def __init__(self, device=None, platform=None, browser=None, options=None):
         self.headers: Headers
 
     def __find_device(self) -> str:
-        if self.device is not None:
-            if utils.contains_multiple(self.device, devices):
-                self.device = utils.choice(self.device)
-            else:
-                raise exceptions.InvalidArgumentError('No such device type found: {}'.format(self.device))
+        if self.device is not None and self.device not in devices:
+            raise exceptions.InvalidArgumentError('No such device type found: {}'.format(self.device))
 
         # Override the device type, if the platform is specified
         if self.platform is not None:
-            if utils.contains_multiple(self.platform, platforms_desktop):
+            if self.platform in platforms_desktop:
                 self.device = 'desktop'
-            elif utils.contains_multiple(self.platform, platforms_mobile):
+            elif self.platform in platforms_mobile:
                 self.device = 'mobile'
 
         if self.device is None:
@@ -45,20 +42,17 @@ def __find_device(self) -> str:
         return self.device
 
     def __find_platform(self) -> str:
-        if self.platform is not None:
-            if utils.contains_multiple(self.platform, platforms):
-                self.platform = utils.choice(self.platform)
-            else:
-                raise exceptions.InvalidArgumentError('No such platform found: {}'.format(self.platform))
+        if self.platform is not None and self.platform not in platforms:
+            raise exceptions.InvalidArgumentError('No such platform found: {}'.format(self.platform))
 
         # Make the platform consistent with the device type and browser
-        if self.device == 'desktop' and not utils.contains_multiple(self.platform, platforms_desktop):
+        if self.device == 'desktop' and self.platform not in platforms_desktop:
             # Safari only supports the macOS and iOS platforms
             if self.browser is not None and self.browser == 'safari':
                 self.platform = utils.choice(('macos', 'ios'))
             else:
                 self.platform = utils.choice(platforms_desktop)
-        elif self.device == 'mobile' and not utils.contains_multiple(self.platform, platforms_mobile):
+        elif self.device == 'mobile' and self.platform not in platforms_mobile:
             self.platform = utils.choice(platforms_mobile)
 
         if self.platform is None:
@@ -67,11 +61,8 @@ def __find_platform(self) -> str:
         return self.platform
 
     def __find_browser(self) -> str:
-        if self.browser is not None:
-            if utils.contains_multiple(self.browser, browsers):
-                self.browser = utils.choice(self.browser)
-            else:
-                raise exceptions.InvalidArgumentError('No such browser found: {}'.format(self.browser))
+        if self.browser is not None and self.browser not in browsers:
+            raise exceptions.InvalidArgumentError('No such browser found: {}'.format(self.browser))
 
         if self.browser is None:
             self.browser = utils.choice(browsers)

src/ua_generator/utils.py

@@ -7,24 +7,7 @@
 from typing import Union
 
 
-def contains(t: Union[str, tuple, list], val: str) -> bool:
-    if type(t) is str and t == val:
-        return True
-    if (type(t) is tuple or type(t) is list) and val in t:
-        return True
-
-    return False
-
-
-def contains_multiple(t: Union[str, tuple, list], arr: Union[tuple, list]) -> bool:
-    for val in arr:
-        if contains(t, val):
-            return True
-
-    return False
-
-
-def choice(t):
+def choice(t: Union[str, tuple, list, None]) -> Union[str, None]:
     if type(t) is str:
         return t
     if type(t) is tuple or type(t) is list:

tests/test_exceptions.py

@@ -10,11 +10,30 @@
 
 
 def raised_call():
-    ua_generator.generate(device='desktop', platform='commodore_64')
+    ua_generator.generate(device='desktop', platform='invalid111')
 
 
 def raised_call_2():
-    ua_generator.generate(browser=('netscape', 'ie'))
+    ua_generator.generate(browser=('invalid111', 'invalid112'))
+
+
+def raised_call_3():
+    ua_generator.generate(device='invalid111', platform='android', browser='chrome')
+
+
+def raised_call_4():
+    for i in range(0, 100):
+        ua_generator.generate(device=('desktop', 'invalid111'))
+
+
+def raised_call_5():
+    for i in range(0, 100):
+        ua_generator.generate(platform=('invalid111', 'macos'))
+
+
+def raised_call_6():
+    for i in range(0, 100):
+        ua_generator.generate(browser=('invalid111', 'chrome'))
 
 
 class TestExceptions(unittest.TestCase):
@@ -24,6 +43,18 @@ def test_value_error(self):
     def test_value_error_2(self):
         self.assertRaises(exceptions.InvalidArgumentError, raised_call_2)
 
+    def test_value_error_3(self):
+        self.assertRaises(exceptions.InvalidArgumentError, raised_call_3)
+
+    def test_value_error_4(self):
+        self.assertRaises(exceptions.InvalidArgumentError, raised_call_4)
+
+    def test_value_error_5(self):
+        self.assertRaises(exceptions.InvalidArgumentError, raised_call_5)
+
+    def test_value_error_6(self):
+        self.assertRaises(exceptions.InvalidArgumentError, raised_call_6)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_utils.py

@@ -9,21 +9,12 @@
 
 
 class TestUtils(unittest.TestCase):
-    def test_contains(self):
-        self.assertTrue(utils.contains(('mobile', 'desktop'), 'mobile'))
-        self.assertTrue(utils.contains(('mobile', 'desktop'), 'desktop'))
-        self.assertTrue(utils.contains('desktop', 'desktop'))
-        self.assertFalse(utils.contains(('mobile', 'desktop'), 'ebook'))
-        self.assertFalse(utils.contains('desktop', 'ebook'))
-
-    def test_contains_multiple(self):
-        self.assertTrue(utils.contains_multiple(('mobile', 'desktop'), ['mobile', 'desktop']))
-        self.assertTrue(utils.contains_multiple(('mobile', 'desktop'), ['ebook', 'desktop']))
-        self.assertTrue(utils.contains_multiple(['mobile', 'desktop'], ['ebook', 'mobile']))
-        self.assertTrue(utils.contains_multiple('desktop', ('ebook', 'desktop')))
-        self.assertTrue(utils.contains_multiple('desktop', ('ebook', 'desktop')))
-        self.assertFalse(utils.contains_multiple(['mobile', 'desktop'], ['ebook', 'gameboy']))
-        self.assertFalse(utils.contains_multiple('mobile', ['ebook', 'gameboy']))
+    def test_choice(self):
+        self.assertEqual(utils.choice('mobile'), 'mobile')
+        self.assertEqual(utils.choice(('mobile', 'mobile')), 'mobile')
+        self.assertEqual(utils.choice(['mobile', 'mobile']), 'mobile')
+        self.assertEqual(utils.choice(None), None)
+        self.assertEqual(utils.choice(111), None)
 
 
 if __name__ == '__main__':