test_utils.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. import atheris # pragma: no cover
  2. @atheris.instrument_func
  3. def is_expected_exception(
  4. error_message_list: list[str], exception: Exception
  5. ): # pragma: no cover
  6. """Checks if the message of a given exception matches any of the expected error messages.
  7. Args:
  8. error_message_list (List[str]): A list of error message substrings to check against the exception's message.
  9. exception (Exception): The exception object raised during execution.
  10. Returns:
  11. bool: True if the exception's message contains any of the substrings from the error_message_list, otherwise False.
  12. """
  13. for error in error_message_list:
  14. if error in str(exception):
  15. return True
  16. return False
  17. class EnhancedFuzzedDataProvider(atheris.FuzzedDataProvider): # pragma: no cover
  18. """Extends atheris.FuzzedDataProvider to offer additional methods to make fuzz testing slightly more DRY."""
  19. def __init__(self, data):
  20. """Initializes the EnhancedFuzzedDataProvider with fuzzing data from the argument provided to TestOneInput.
  21. Args:
  22. data (bytes): The binary data used for fuzzing.
  23. """
  24. super().__init__(data)
  25. def ConsumeRemainingBytes(self) -> bytes:
  26. """Consume the remaining bytes in the bytes container.
  27. Returns:
  28. bytes: Zero or more bytes.
  29. """
  30. return self.ConsumeBytes(self.remaining_bytes())
  31. def ConsumeRandomBytes(self, max_length=None) -> bytes:
  32. """Consume a random count of bytes from the bytes container.
  33. Args:
  34. max_length (int, optional): The maximum length of the string. Defaults to the number of remaining bytes.
  35. Returns:
  36. bytes: Zero or more bytes.
  37. """
  38. if max_length is None:
  39. max_length = self.remaining_bytes()
  40. else:
  41. max_length = min(max_length, self.remaining_bytes())
  42. return self.ConsumeBytes(self.ConsumeIntInRange(0, max_length))
  43. def ConsumeRandomString(self, max_length=None, without_surrogates=False) -> str:
  44. """Consume bytes to produce a Unicode string.
  45. Args:
  46. max_length (int, optional): The maximum length of the string. Defaults to the number of remaining bytes.
  47. without_surrogates (bool, optional): If True, never generate surrogate pair characters. Defaults to False.
  48. Returns:
  49. str: A Unicode string.
  50. """
  51. if max_length is None:
  52. max_length = self.remaining_bytes()
  53. else:
  54. max_length = min(max_length, self.remaining_bytes())
  55. count = self.ConsumeIntInRange(0, max_length)
  56. if without_surrogates:
  57. return self.ConsumeUnicodeNoSurrogates(count)
  58. else:
  59. return self.ConsumeUnicode(count)
  60. def ConsumeRandomInt(self, minimum=0, maximum=1234567890) -> int:
  61. """Consume bytes to produce an integer.
  62. Args:
  63. minimum (int, optional): The minimum value of the integer. Defaults to 0.
  64. maximum (int, optional): The maximum value of the integer. Defaults to 1234567890.
  65. Returns:
  66. int: An integer.
  67. """
  68. return self.ConsumeIntInRange(minimum, maximum)