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 24from __future__ import (absolute_import, division, print_function, 25 unicode_literals) 26 27import logging 28import os 29import re 30 31log = logging.getLogger(__name__) 32 33 34REPLY_DATA = re.compile("<reply>[ \t\n\r]*<data[^<]*>(.*?)</data>", re.MULTILINE | re.DOTALL) 35 36 37class ClosingFileHandler(logging.StreamHandler): 38 def __init__(self, filename): 39 super(ClosingFileHandler, self).__init__() 40 self.filename = os.path.abspath(filename) 41 self.setStream(None) 42 43 def emit(self, record): 44 with open(self.filename, "a") as fp: 45 self.setStream(fp) 46 super(ClosingFileHandler, self).emit(record) 47 self.setStream(None) 48 49 def setStream(self, stream): 50 setStream = getattr(super(ClosingFileHandler, self), 'setStream', None) 51 if callable(setStream): 52 return setStream(stream) 53 if stream is self.stream: 54 result = None 55 else: 56 result = self.stream 57 self.acquire() 58 try: 59 self.flush() 60 self.stream = stream 61 finally: 62 self.release() 63 return result 64 65class TestData(object): 66 def __init__(self, data_folder): 67 self.data_folder = data_folder 68 69 def get_test_data(self, test_number): 70 # Create the test file name 71 filename = os.path.join(self.data_folder, 72 "test{0}".format(test_number)) 73 74 log.debug("Parsing file %s", filename) 75 76 with open(filename, "rb") as f: 77 contents = f.read().decode("utf-8") 78 79 m = REPLY_DATA.search(contents) 80 if not m: 81 raise Exception("Couldn't find a <reply><data> section") 82 83 # Left-strip the data so we don't get a newline before our data. 84 return m.group(1).lstrip() 85 86 87if __name__ == '__main__': 88 td = TestData("./data") 89 data = td.get_test_data(1) 90 print(data) 91