diff --git a/src/ua_generator/user_agent.py b/src/ua_generator/user_agent.py index 82ddb4f..22b4e77 100644 --- a/src/ua_generator/user_agent.py +++ b/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) diff --git a/src/ua_generator/utils.py b/src/ua_generator/utils.py index a299db9..07d3ea8 100644 --- a/src/ua_generator/utils.py +++ b/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: diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py index b2a07fd..caa5ea5 100644 --- a/tests/test_exceptions.py +++ b/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() diff --git a/tests/test_utils.py b/tests/test_utils.py index 694cdde..f757820 100644 --- a/tests/test_utils.py +++ b/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__':