Best Python code snippet using yandex-tank
missile.py
Source:missile.py  
1'''2Missile object and generators3You should update Stepper.status.ammo_count and Stepper.status.loop_count in your custom generators!4'''5import logging6from itertools import cycle7from netort.resource import manager as resource8from . import info9from .module_exceptions import AmmoFileError10class HttpAmmo(object):11    '''12    Represents HTTP missile13    >>> print HttpAmmo('/', []).to_s()  # doctest: +NORMALIZE_WHITESPACE14    GET / HTTP/1.115    >>> print HttpAmmo('/', ['Connection: Close', 'Content-Type: Application/JSON']).to_s()  # doctest: +NORMALIZE_WHITESPACE16    GET / HTTP/1.117    Connection: Close18    Content-Type: Application/JSON19    >>> print HttpAmmo('/', ['Connection: Close'], method='POST', body='hello!').to_s()  # doctest: +NORMALIZE_WHITESPACE20    POST / HTTP/1.121    Connection: Close22    Content-Length: 623    <BLANKLINE>24    hello!25    '''26    def __init__(self, uri, headers, method='GET', http_ver='1.1', body=b''):27        self.method = method28        self.uri = uri.encode('utf8') if isinstance(uri, str) else uri29        self.proto = 'HTTP/%s' % http_ver30        self.headers = set(headers)31        self.body = body32        if len(body):33            self.headers.add("Content-Length: %s" % len(body))34    def to_s(self):35        if self.headers:36            headers = b'\r\n'.join(sorted([h.encode('utf8') for h in self.headers])) + b'\r\n'37        else:38            headers = b''39        return b"%s %s %s\r\n%s\r\n%s" % (40            self.method.encode('utf8'),41            self.uri,42            self.proto.encode('utf8'),43            headers,44            self.body)45class SimpleGenerator(object):46    '''47    Generates ammo based on a given sample.48    '''49    def __init__(self, missile_sample):50        '''51        Missile sample is any object that has to_s method which52        returns its string representation.53        '''54        self.missiles = cycle([(missile_sample.to_s(), None)])55    def __iter__(self):56        for m in self.missiles:57            info.status.inc_loop_count()58            yield m59class UriStyleGenerator(object):60    '''61    Generates GET ammo based on given URI list.62    '''63    def __init__(self, uris, headers, http_ver='1.1'):64        '''65        uris - a list of URIs as strings.66        '''67        self.uri_count = len(uris)68        self.missiles = cycle([(69            HttpAmmo(70                uri, headers, http_ver=http_ver).to_s(), None) for uri in uris])71    def __iter__(self):72        for m in self.missiles:73            yield m74            info.status.loop_count = info.status.ammo_count / self.uri_count75class Reader(object):76    def __init__(self, filename, use_cache=True, **kwargs):77        self.filename = filename78        self.use_cache = use_cache79class AmmoFileReader(Reader):80    '''Read missiles from ammo file'''81    def __init__(self, filename, use_cache=True, **kwargs):82        super(AmmoFileReader, self).__init__(filename, use_cache)83        self.log = logging.getLogger(__name__)84        self.log.info("Loading ammo from '%s'" % filename)85    def __iter__(self):86        def read_chunk_header(ammo_file):87            chunk_header = b''88            while chunk_header == b'':89                line = ammo_file.readline()90                if line == b'':91                    return line92                chunk_header = line.strip(b'\r\n')93            return chunk_header94        opener = resource.get_opener(self.filename)95        with opener(self.use_cache) as ammo_file:96            info.status.af_size = opener.data_length97            # if we got StopIteration here, the file is empty98            chunk_header = read_chunk_header(ammo_file)99            while chunk_header:100                if chunk_header != b'':101                    try:102                        fields = chunk_header.split()103                        chunk_size = int(fields[0])104                        if chunk_size == 0:105                            if info.status.loop_count == 0:106                                self.log.info(107                                    'Zero-sized chunk in ammo file at %s. Starting over.'108                                    % ammo_file.tell())109                            ammo_file.seek(0)110                            info.status.inc_loop_count()111                            chunk_header = read_chunk_header(ammo_file)112                            continue113                        marker = fields[1] if len(fields) > 1 else None114                        missile = ammo_file.read(chunk_size)115                        if len(missile) < chunk_size:116                            raise AmmoFileError(117                                "Unexpected end of file: read %s bytes instead of %s"118                                % (len(missile), chunk_size))119                        yield (missile, marker)120                    except (IndexError, ValueError) as e:121                        raise AmmoFileError(122                            "Error while reading ammo file. Position: %s, header: '%s', original exception: %s"123                            % (ammo_file.tell(), chunk_header, e))124                chunk_header = read_chunk_header(ammo_file)125                if chunk_header == b'':126                    ammo_file.seek(0)127                    info.status.inc_loop_count()128                    chunk_header = read_chunk_header(ammo_file)129                info.status.af_position = ammo_file.tell()130class SlowLogReader(Reader):131    """Read missiles from SQL slow log. Not usable with Phantom"""132    def __iter__(self):133        opener = resource.get_opener(self.filename)134        with opener(self.use_cache) as ammo_file:135            info.status.af_size = opener.data_length136            request = ""137            while True:138                for line in ammo_file:139                    info.status.af_position = ammo_file.tell()140                    if isinstance(line, bytes):141                        line = line.decode('utf-8')142                    if line.startswith('#'):143                        if request != "":144                            yield (request, None)145                            request = ""146                    else:147                        request += line148                ammo_file.seek(0)149                info.status.af_position = 0150                info.status.inc_loop_count()151class LineReader(Reader):152    """One line -- one missile"""153    def __iter__(self):154        opener = resource.get_opener(self.filename)155        with opener(self.use_cache) as ammo_file:156            info.status.af_size = opener.data_length157            while True:158                for line in ammo_file:159                    info.status.af_position = ammo_file.tell()160                    yield (line.rstrip(b'\r\n'), None) if isinstance(line, bytes) else (line.rstrip('\r\n').encode('utf8'), None)161                ammo_file.seek(0)162                info.status.af_position = 0163                info.status.inc_loop_count()164class CaseLineReader(Reader):165    """One line -- one missile with case, tab separated"""166    def __iter__(self):167        opener = resource.get_opener(self.filename)168        with opener(self.use_cache) as ammo_file:169            info.status.af_size = opener.data_length170            while True:171                for line in ammo_file:172                    info.status.af_position = ammo_file.tell()173                    parts = line.rstrip(b'\r\n').split(b'\t', 1)174                    if len(parts) == 2:175                        yield (parts[1], parts[0])176                    elif len(parts) == 1:177                        yield (parts[0], None)178                    else:179                        raise RuntimeError("Unreachable branch")180                ammo_file.seek(0)181                info.status.af_position = 0182                info.status.inc_loop_count()183class AccessLogReader(Reader):184    """Missiles from access log"""185    def __init__(self, filename, headers=None, http_ver='1.1', use_cache=True, **kwargs):186        super(AccessLogReader, self).__init__(filename, use_cache)187        self.warned = False188        self.headers = set(headers) if headers else set()189        self.log = logging.getLogger(__name__)190    def warn(self, message):191        if not self.warned:192            self.warned = True193            self.log.warning(194                "There are some skipped lines. See full log for details.")195        self.log.debug(message)196    def __iter__(self):197        opener = resource.get_opener(self.filename)198        with opener(self.use_cache) as ammo_file:199            info.status.af_size = opener.data_length200            while True:201                for line in ammo_file:202                    info.status.af_position = ammo_file.tell()203                    if isinstance(line, bytes):204                        line = line.decode('utf-8')205                    try:206                        request = line.split('"')[1]207                        method, uri, proto = request.split()208                        http_ver = proto.split('/')[1]209                        if method == "GET":210                            yield (211                                HttpAmmo(212                                    uri,213                                    headers=self.headers,214                                    http_ver=http_ver, ).to_s(), None)215                        else:216                            self.warn(217                                "Skipped line: %s (unsupported method)" % line)218                    except (ValueError, IndexError) as e:219                        self.warn("Skipped line: %s (%s)" % (line, e))220                ammo_file.seek(0)221                info.status.af_position = 0222                info.status.inc_loop_count()223def _parse_header(header):224    return dict([(h.strip().decode('utf8') for h in header.split(b':', 1))])225class UriReader(Reader):226    def __init__(self, filename, headers=None, http_ver='1.1', use_cache=True, **kwargs):227        super(UriReader, self).__init__(filename, use_cache)228        self.headers = {pair[0].strip(): pair[1].strip() for pair in [h.split(':', 1) for h in headers]} \229            if headers else {}230        self.http_ver = http_ver231        self.log = logging.getLogger(__name__)232        self.log.info("Loading ammo from '%s' using URI format." % filename)233    def __iter__(self):234        opener = resource.get_opener(self.filename)235        with opener(self.use_cache) as ammo_file:236            info.status.af_size = opener.data_length237            while True:238                for line in ammo_file:239                    info.status.af_position = ammo_file.tell()240                    # if isinstance(line, bytes):241                    #     line = line.decode('utf-8')242                    if line.startswith(b'['):243                        self.headers.update(244                            _parse_header(line.strip(b'\r\n[]\t ')))245                    elif len(line.rstrip(b'\r\n')):246                        fields = line.split()247                        uri = fields[0]248                        if len(fields) > 1:249                            marker = fields[1]250                        else:251                            marker = None252                        yield (253                            HttpAmmo(254                                uri,255                                headers=[256                                    ': '.join(header)257                                    for header in self.headers.items()258                                ],259                                http_ver=self.http_ver, ).to_s(), marker)260                if info.status.ammo_count == 0:261                    self.log.error("No ammo in uri-style file")262                    raise AmmoFileError("No ammo! Cover me!")263                ammo_file.seek(0)264                info.status.af_position = 0265                info.status.inc_loop_count()266class UriPostReader(Reader):267    """Read POST missiles from ammo file"""268    def __init__(self, filename, headers=None, http_ver='1.1', use_cache=True, **kwargs):269        super(UriPostReader, self).__init__(filename, use_cache)270        self.headers = {pair[0].strip(): pair[1].strip() for pair in [h.split(':', 1) for h in headers]} \271            if headers else {}272        self.http_ver = http_ver273        self.log = logging.getLogger(__name__)274        self.log.info("Loading ammo from '%s' using URI+POST format", filename)275    def __iter__(self):276        def read_chunk_header(ammo_file):277            chunk_header = b''278            while chunk_header == b'':279                line = ammo_file.readline()280                if line.startswith(b'['):281                    self.headers.update(_parse_header(line.strip(b'\r\n[]\t ')))282                elif line == b'':283                    return line284                else:285                    chunk_header = line.strip(b'\r\n')286            return chunk_header287        opener = resource.get_opener(self.filename)288        with opener(self.use_cache) as ammo_file:289            info.status.af_size = opener.data_length290            # if we got StopIteration here, the file is empty291            chunk_header = read_chunk_header(ammo_file)292            while chunk_header:293                if chunk_header != b'':294                    try:295                        fields = chunk_header.split()296                        chunk_size = int(fields[0])297                        uri = fields[1]298                        marker = fields[2] if len(fields) > 2 else None299                        if chunk_size == 0:300                            missile = b""301                        else:302                            missile = ammo_file.read(chunk_size)303                        if len(missile) < chunk_size:304                            raise AmmoFileError(305                                "Unexpected end of file: read %s bytes instead of %s"306                                % (len(missile), chunk_size))307                        yield (308                            HttpAmmo(309                                uri=uri,310                                headers=[311                                    ': '.join(header)312                                    for header in self.headers.items()313                                ],314                                method='POST',315                                body=missile,316                                http_ver=self.http_ver, ).to_s(), marker)317                    except (IndexError, ValueError) as e:318                        raise AmmoFileError(319                            "Error while reading ammo file. Position: %s, header: '%s', original exception: %s"320                            % (ammo_file.tell(), chunk_header, e))321                chunk_header = read_chunk_header(ammo_file)322                if chunk_header == '':323                    self.log.debug(324                        'Reached the end of ammo file. Starting over.')325                    ammo_file.seek(0)326                    info.status.inc_loop_count()327                    chunk_header = read_chunk_header(ammo_file)...info.py
Source:info.py  
...36                "Tried to publish to a non-existent key: %s" % key)37        log.debug('Published %s to %s', value, key)38        self.info[key] = value39    @property40    def af_position(self):41        return self._af_position42    @af_position.setter43    def af_position(self, value):44        self._af_position = value45        self.update_af_progress()46    @property47    def ammo_count(self):48        return self._ammo_count49    @ammo_count.setter50    def ammo_count(self, value):51        self._ammo_count = value52        self.update_lp_progress()53        if self.ammo_limit and value > self.ammo_limit:54            print()55            log.info("Ammo limit reached: %s", self.ammo_limit)56            raise StopIteration57    def inc_ammo_count(self):...Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!
