util.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. #
  4. # Project ___| | | | _ \| |
  5. # / __| | | | |_) | |
  6. # | (__| |_| | _ <| |___
  7. # \___|\___/|_| \_\_____|
  8. #
  9. # Copyright (C) 2017 - 2020, Daniel Stenberg, <daniel@haxx.se>, et al.
  10. #
  11. # This software is licensed as described in the file COPYING, which
  12. # you should have received as part of this distribution. The terms
  13. # are also available at https://curl.se/docs/copyright.html.
  14. #
  15. # You may opt to use, copy, modify, merge, publish, distribute and/or sell
  16. # copies of the Software, and permit persons to whom the Software is
  17. # furnished to do so, under the terms of the COPYING file.
  18. #
  19. # This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY
  20. # KIND, either express or implied.
  21. #
  22. """Module for extracting test data from the test data folder and other utils"""
  23. from __future__ import (absolute_import, division, print_function,
  24. unicode_literals)
  25. import logging
  26. import os
  27. import re
  28. log = logging.getLogger(__name__)
  29. REPLY_DATA = re.compile("<reply>[ \t\n\r]*<data[^<]*>(.*?)</data>", re.MULTILINE | re.DOTALL)
  30. class ClosingFileHandler(logging.StreamHandler):
  31. def __init__(self, filename):
  32. super(ClosingFileHandler, self).__init__()
  33. self.filename = os.path.abspath(filename)
  34. self.setStream(None)
  35. def emit(self, record):
  36. with open(self.filename, "a") as fp:
  37. self.setStream(fp)
  38. super(ClosingFileHandler, self).emit(record)
  39. self.setStream(None)
  40. def setStream(self, stream):
  41. setStream = getattr(super(ClosingFileHandler, self), 'setStream', None)
  42. if callable(setStream):
  43. return setStream(stream)
  44. if stream is self.stream:
  45. result = None
  46. else:
  47. result = self.stream
  48. self.acquire()
  49. try:
  50. self.flush()
  51. self.stream = stream
  52. finally:
  53. self.release()
  54. return result
  55. class TestData(object):
  56. def __init__(self, data_folder):
  57. self.data_folder = data_folder
  58. def get_test_data(self, test_number):
  59. # Create the test file name
  60. filename = os.path.join(self.data_folder,
  61. "test{0}".format(test_number))
  62. log.debug("Parsing file %s", filename)
  63. with open(filename, "rb") as f:
  64. contents = f.read().decode("utf-8")
  65. m = REPLY_DATA.search(contents)
  66. if not m:
  67. raise Exception("Couldn't find a <reply><data> section")
  68. # Left-strip the data so we don't get a newline before our data.
  69. return m.group(1).lstrip()
  70. if __name__ == '__main__':
  71. td = TestData("./data")
  72. data = td.get_test_data(1)
  73. print(data)