feat: ext

pyresttest/ext/extracotr_cookie.py

@@ -0,0 +1,53 @@
+# coding=utf-8
+
+import sys
+
+import pyresttest.validators as validators
+from pyresttest.binding import Context
+
+# Python 3 compatibility
+if sys.version_info[0] > 2:
+    from past.builtins import basestring
+from pyresttest.six import text_type
+from pyresttest.six import binary_type
+
+def _first(lst):
+    if len(lst) > 0:
+        return lst[0]
+    return None
+
+
+def _list_to_multi_value_map(lst):
+    a_map = {}
+    for k, v in lst:
+        if k not in a_map:
+            a_map[k] = v
+        else:
+            if isinstance(a_map[k], list):
+                a_map[k].append(v)
+            else:
+                a_map[k] = [a_map[k], v]
+    return a_map
+
+
+class CookieExtractor(validators.AbstractExtractor):
+    """get a wanted cookie"""
+
+    extractor_type = 'cookie'
+    is_body_extractor = True
+    HEADER_NAME = 'Set-Cookie'
+
+    @classmethod
+    def parse(cls, config, extractor_base=None):
+        base = CookieExtractor()
+        return cls.configure_base(config, base)
+
+    def extract_internal(self, query=None, args=None, body=None, headers=None):
+        headers_map = _list_to_multi_value_map(headers)
+        return _first([
+            x for x in headers_map.get(CookieExtractor.HEADER_NAME.lower(), [])
+            if ('%s=' % query) in x
+        ])
+
+
+EXTRACTORS = {'cookie': CookieExtractor.parse}

pyresttest/ext/file_choice_generator.py

@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+
+# file choice generator
+
+import random
+import linecache
+import sys
+
+# Python 3 compatibility
+if sys.version_info[0] > 2:
+    from past.builtins import basestring
+
+
+def factory_file_choice_generator(values):
+    """ Returns generators that picks values from a certain file"""
+
+    def file_choice_generator():
+        with open(values, 'rU') as my_file:
+            cnt_list = my_file.readlines()  #list+choice的方式,如果读取文件较大,可能会导致list占用较多内不能,优化可改成yield计算count
+            # count = len(my_file.readlines())
+            if len(cnt_list) == 0:
+                raise ValueError('The length of file %s is null' %(values))
+            while(True):
+                # cnt = linecache.getline(my_file, rand)
+                yield random.choice(cnt_list).strip()
+    return file_choice_generator
+
+def parse_file_choice_generator(config):
+    """ Parse file choice generator """
+    vals = config['values']
+    if not vals:
+        raise ValueError('Values for choice filepath must exist')
+    if not isinstance(vals, basestring):
+        raise ValueError('Values must be a basestring')
+    with open(vals) as f:
+        if not isinstance(f, file):
+            raise ValueError('File opened from filepath is not valid')
+    return factory_file_choice_generator(vals)()
+
+
+
+# This is where the magic happens, each one of these is a registry of
+# validators/extractors/generators to use
+# VALIDATORS = {'contains': ContainsValidator.parse}
+# VALIDATOR_TESTS = {'is_dict': test_is_dict}
+
+# Converts to lowercase and tests for equality
+# COMPARATORS = {'str.eq.lower': lambda a, b: str(a).lower() == str(b).lower()}
+
+# EXTRACTORS = {'weirdzo': WeirdzoExtractor.parse}
+GENERATORS = {'file_choice': parse_file_choice_generator}

pyresttest/ext/file_seq_generator.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+# file sequence generator
+
+import sys
+import os
+
+# Python 3 compatibility
+if sys.version_info[0] > 2:
+    from past.builtins import basestring
+
+
+def factory_file_seq_generator(values):
+    """ Returns generators that picks values from a certain file in sequence"""
+
+    def file_seq_generator():
+        with open(values, 'rU') as my_file:
+            cnt_list = my_file.readlines()
+            i = 0
+            while(True):
+                yield cnt_list[i].strip()
+                if i == len(cnt_list):
+                    i == 0
+    return file_seq_generator
+
+def parse_file_seq_generator(config):
+    """ Parse file seq generator """
+    vals = config['values']
+    if not vals:
+        raise ValueError('Values for choice filepath must exist')
+    if not isinstance(vals, basestring):
+        raise ValueError('Values must be a basestring')
+    # with open(vals) as f:
+        # if not isinstance(f, file):
+    if not os.path.isfile(vals):
+        raise ValueError('File opened from filepath is not valid')
+    return factory_file_seq_generator(vals)()
+
+
+
+# This is where the magic happens, each one of these is a registry of
+# validators/extractors/generators to use
+# VALIDATORS = {'contains': ContainsValidator.parse}
+# VALIDATOR_TESTS = {'is_dict': test_is_dict}
+
+# Converts to lowercase and tests for equality
+# COMPARATORS = {'str.eq.lower': lambda a, b: str(a).lower() == str(b).lower()}
+
+# EXTRACTORS = {'weirdzo': WeirdzoExtractor.parse}
+GENERATORS = {'file_seq': parse_file_seq_generator}