consistency

src/ua_generator/user_agent.py

@@ -25,52 +25,59 @@ def __init__(self, device=None, platform=None, browser=None, options=None):
         self.ch: ClientHints
         self.headers: Headers
 
-    def __find_device(self):
+    def __find_device(self) -> str:
         if self.device is not None:
-            if not utils.contains_multiple(self.device, devices):
-                raise exceptions.InvalidArgumentError('No device type found: {}'.format(self.device))
-            else:
+            if utils.contains_multiple(self.device, devices):
                 self.device = utils.choice(self.device)
-                return self.device
+            else:
+                raise exceptions.InvalidArgumentError('No such device type found: {}'.format(self.device))
 
-        if self.platform is not None and utils.contains_multiple(self.platform, platforms_desktop):
-            self.device = 'desktop'
-        elif self.platform is not None and utils.contains_multiple(self.platform, platforms_mobile):
-            self.device = 'mobile'
-        elif self.device is None:
+        # Override the device type, if the platform is specified
+        if self.platform is not None:
+            if utils.contains_multiple(self.platform, platforms_desktop):
+                self.device = 'desktop'
+            elif utils.contains_multiple(self.platform, platforms_mobile):
+                self.device = 'mobile'
+
+        if self.device is None:
             self.device = utils.choice(devices)
 
         return self.device
 
-    def __find_platform(self):
+    def __find_platform(self) -> str:
         if self.platform is not None:
-            if not utils.contains_multiple(self.platform, platforms):
-                raise exceptions.InvalidArgumentError('No platform found: {}'.format(self.platform))
-            else:
+            if utils.contains_multiple(self.platform, platforms):
                 self.platform = utils.choice(self.platform)
-                return self.platform
+            else:
+                raise exceptions.InvalidArgumentError('No such platform found: {}'.format(self.platform))
 
-        if self.device is not None and 'desktop' in self.device:
-            self.platform = utils.choice(platforms_desktop)
-        elif self.device is not None and 'mobile' in self.device:
+        # Make the platform consistent with the device type and browser
+        if self.device == 'desktop' and not utils.contains_multiple(self.platform, 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):
             self.platform = utils.choice(platforms_mobile)
-        elif self.platform is None:
+
+        if self.platform is None:
             self.platform = utils.choice(platforms)
 
         return self.platform
 
-    def __find_browser(self):
+    def __find_browser(self) -> str:
         if self.browser is not None:
-            if not utils.contains_multiple(self.browser, browsers):
-                raise exceptions.InvalidArgumentError('No browser found: {}'.format(self.browser))
-            else:
+            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 None:
             self.browser = utils.choice(browsers)
 
-        # Safari only support for macOS and iOS
-        if self.platform != 'macos' and self.platform != 'ios' and self.browser == 'safari':
+        # Safari only supports the macOS and iOS platforms
+        if self.browser == 'safari' and self.platform not in ('macos', 'ios'):
             self.browser = 'chrome'
 
         return self.browser

src/ua_generator/utils.py

@@ -4,9 +4,10 @@
 License: Apache License 2.0
 """
 import random
+from typing import Union
 
 
-def contains(t, val):
+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:
@@ -15,7 +16,7 @@ def contains(t, val):
     return False
 
 
-def contains_multiple(t, arr):
+def contains_multiple(t: Union[str, tuple, list], arr: Union[tuple, list]) -> bool:
     for val in arr:
         if contains(t, val):
             return True
@@ -24,11 +25,9 @@ def contains_multiple(t, arr):
 
 
 def choice(t):
-    if not t:
-        return None
     if type(t) is str:
         return t
     if type(t) is tuple or type(t) is list:
         return random.choice(t)
 
-    return False
+    return None

tests/test_browser.py

@@ -11,7 +11,7 @@
 class TestBrowser(unittest.TestCase):
     def test_browser(self):
         for i in range(0, 100):
-            ua = ua_generator.generate(browser=('chrome'))
+            ua = ua_generator.generate(browser=('chrome',))
             self.assertTrue(ua.browser == 'chrome')
 
     def test_browser_2(self):

tests/test_platform.py

@@ -11,7 +11,7 @@
 class TestPlatform(unittest.TestCase):
     def test_platform(self):
         for i in range(0, 100):
-            ua = ua_generator.generate(platform=('linux'))
+            ua = ua_generator.generate(platform=('linux',))
             self.assertTrue(ua.platform == 'linux')
 
     def test_platform_2(self):
@@ -34,6 +34,30 @@ def test_platform_5(self):
             ua = ua_generator.generate(device='desktop')
             self.assertTrue(ua.platform == 'windows' or ua.platform == 'macos' or ua.platform == 'linux')
 
+    def test_platform_6(self):
+        for i in range(0, 100):
+            ua = ua_generator.generate(device='desktop', platform='ios')
+            self.assertNotEqual(ua.device, 'desktop')
+            self.assertEqual(ua.platform, 'ios')
+
+    def test_platform_7(self):
+        for i in range(0, 100):
+            ua = ua_generator.generate(device='desktop', platform='android')
+            self.assertNotEqual(ua.device, 'desktop')
+            self.assertEqual(ua.platform, 'android')
+
+    def test_platform_8(self):
+        for i in range(0, 100):
+            ua = ua_generator.generate(device='mobile', platform='windows')
+            self.assertNotEqual(ua.device, 'mobile')
+            self.assertEqual(ua.platform, 'windows')
+
+    def test_platform_9(self):
+        for i in range(0, 100):
+            ua = ua_generator.generate(device='mobile', platform='macos')
+            self.assertNotEqual(ua.device, 'mobile')
+            self.assertEqual(ua.platform, 'macos')
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_user_agent.py

@@ -42,6 +42,17 @@ def test_user_agent_4(self):
             self.assertIsNotNone(ua.browser)
             self.assertIsNotNone(ua.text)
 
+    def test_user_agent_5(self):
+        for i in range(0, 1000):
+            device = ('desktop', 'mobile')
+            platform = ('windows', 'macos', 'ios', 'linux', 'android')
+            browser = ('chrome', 'edge', 'firefox', 'safari')
+            ua = ua_generator.generate(device=device, browser=browser, platform=platform)
+            self.assertIsNotNone(ua.device)
+            self.assertIsNotNone(ua.platform)
+            self.assertIsNotNone(ua.browser)
+            self.assertIsNotNone(ua.text)
+
     def test_user_agent_not_contains_brackets(self):
         brackets = re.compile('{(d|s|v|build|chrome|firefox|safari|webkit|windows|android|ios|macos|linux)}')
         for i in range(0, 200):