2
0

util.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. #
  4. # Project ___| | | | _ \| |
  5. # / __| | | | |_) | |
  6. # | (__| |_| | _ <| |___
  7. # \___|\___/|_| \_\_____|
  8. #
  9. # Copyright (C) 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. # SPDX-License-Identifier: curl
  23. #
  24. """Module for extracting test data from the test data folder and other utils"""
  25. from __future__ import (absolute_import, division, print_function,
  26. unicode_literals)
  27. import logging
  28. import os
  29. import re
  30. log = logging.getLogger(__name__)
  31. REPLY_DATA = re.compile("<reply>[ \t\n\r]*<data[^<]*>(.*?)</data>", re.MULTILINE | re.DOTALL)
  32. class ClosingFileHandler(logging.StreamHandler):
  33. def __init__(self, filename):
  34. super(ClosingFileHandler, self).__init__()
  35. self.filename = os.path.abspath(filename)
  36. self.setStream(None)
  37. def emit(self, record):
  38. with open(self.filename, "a") as fp:
  39. self.setStream(fp)
  40. super(ClosingFileHandler, self).emit(record)
  41. self.setStream(None)
  42. def setStream(self, stream):
  43. setStream = getattr(super(ClosingFileHandler, self), 'setStream', None)
  44. if callable(setStream):
  45. return setStream(stream)
  46. if stream is self.stream:
  47. result = None
  48. else:
  49. result = self.stream
  50. self.acquire()
  51. try:
  52. self.flush()
  53. self.stream = stream
  54. finally:
  55. self.release()
  56. return result
  57. class TestData(object):
  58. def __init__(self, data_folder):
  59. self.data_folder = data_folder
  60. def get_test_data(self, test_number):
  61. # Create the test file name
  62. filename = os.path.join(self.data_folder,
  63. "test{0}".format(test_number))
  64. log.debug("Parsing file %s", filename)
  65. with open(filename, "rb") as f:
  66. contents = f.read().decode("utf-8")
  67. m = REPLY_DATA.search(contents)
  68. if not m:
  69. raise Exception("Couldn't find a <reply><data> section")
  70. # Left-strip the data so we don't get a newline before our data.
  71. return m.group(1).lstrip()
  72. if __name__ == '__main__':
  73. td = TestData("./data")
  74. data = td.get_test_data(1)
  75. print(data)