Developer friendly load testing framework Based on asyncio
# encoding: utf-8

from aiotest.user import User, AsyncHttpUser
from aiotest.shape import LoadUserShape
from aiotest import events
from aiotest.client import ClientSession, ClientTimeout
from aiotest.log import logger


__version__ = "0.5.8"

__all__ = (
    "User", 
    "AsyncHttpUser", 
    "LoadUserShape", 
    "events", 
    "ClientSession", 
    "ClientTimeout",
    "logger"
)
aiotest/__init__.py
hewei198711/Aiotest
# encoding: utf8

import sys
import asyncio
from aiotest.main import run


run()
aiotest/__main__.py
hewei198711/Aiotest
# encoding: utf-8

from timeit import default_timer
import asyncio
import base64
import hashlib
import json
import os
import sys
import traceback
import warnings
from contextlib import suppress
from timeit import default_timer
from types import SimpleNamespace, TracebackType
from typing import (
    Any,
    Awaitable,
    Callable,
    Coroutine,
    FrozenSet,
    Generator,
    Generic,
    Iterable,
    List,
    Mapping,
    Optional,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import attr
from multidict import CIMultiDict, MultiDict, MultiDictProxy, istr
from yarl import URL

from aiohttp import hdrs, http, payload
from aiohttp.abc import AbstractCookieJar
from aiohttp.client_exceptions import (
    ClientConnectionError as ClientConnectionError,
    ClientConnectorCertificateError as ClientConnectorCertificateError,
    ClientConnectorError as ClientConnectorError,
    ClientConnectorSSLError as ClientConnectorSSLError,
    ClientError as ClientError,
    ClientHttpProxyError as ClientHttpProxyError,
    ClientOSError as ClientOSError,
    ClientPayloadError as ClientPayloadError,
    ClientProxyConnectionError as ClientProxyConnectionError,
    ClientResponseError as ClientResponseError,
    ClientSSLError as ClientSSLError,
    ContentTypeError as ContentTypeError,
    InvalidURL as InvalidURL,
    ServerConnectionError as ServerConnectionError,
    ServerDisconnectedError as ServerDisconnectedError,
    ServerFingerprintMismatch as ServerFingerprintMismatch,
    ServerTimeoutError as ServerTimeoutError,
    TooManyRedirects as TooManyRedirects,
    WSServerHandshakeError as WSServerHandshakeError,
)
from aiohttp.client_reqrep import (
    ClientRequest as ClientRequest,
    ClientResponse as ClientResponse,
    Fingerprint as Fingerprint,
    RequestInfo as RequestInfo,
    _merge_ssl_params,
)
from aiohttp.client_ws import ClientWebSocketResponse as ClientWebSocketResponse
from aiohttp.connector import (
    BaseConnector as BaseConnector,
    NamedPipeConnector as NamedPipeConnector,
    TCPConnector as TCPConnector,
    UnixConnector as UnixConnector,
)
from aiohttp.cookiejar import CookieJar
from aiohttp.helpers import (
    DEBUG,
    PY_36,
    BasicAuth,
    TimeoutHandle,
    ceil_timeout,
    get_env_proxy_for_url,
    get_running_loop,
    sentinel,
    strip_auth_from_url
)
from aiohttp.http import WS_KEY, HttpVersion, WebSocketReader, WebSocketWriter
from aiohttp.http_websocket import WSHandshakeError, WSMessage, ws_ext_gen, ws_ext_parse
from aiohttp.streams import FlowControlDataQueue
from aiohttp.tracing import Trace, TraceConfig
from aiohttp.typedefs import Final, JSONEncoder, LooseCookies, LooseHeaders, StrOrURL
from aiotest import events
from aiotest.exception import CatchResponseError
from aiotest import runners

__all__ = (
    # client_exceptions
    "ClientConnectionError",
    "ClientConnectorCertificateError",
    "ClientConnectorError",
    "ClientConnectorSSLError",
    "ClientError",
    "ClientHttpProxyError",
    "ClientOSError",
    "ClientPayloadError",
    "ClientProxyConnectionError",
    "ClientResponseError",
    "ClientSSLError",
    "ContentTypeError",
    "InvalidURL",
    "ServerConnectionError",
    "ServerDisconnectedError",
    "ServerFingerprintMismatch",
    "ServerTimeoutError",
    "TooManyRedirects",
    "WSServerHandshakeError",
    # client_reqrep
    "ClientRequest",
    "ClientResponse",
    "Fingerprint",
    "RequestInfo",
    # connector
    "BaseConnector",
    "TCPConnector",
    "UnixConnector",
    "NamedPipeConnector",
    # client_ws
    "ClientWebSocketResponse",
    # client
    "ClientSession",
    "ClientTimeout",
    "request",
)

REQUESTERRORS = (
    ClientResponseError, 
    KeyError, 
    AssertionError, 
    AttributeError, 
    TypeError, 
    IndexError
)

try:
    from ssl import SSLContext
except ImportError:  # pragma: no cover
    SSLContext = object  # type: ignore[misc,assignment]


@attr.s(auto_attribs=True, frozen=True, slots=True)
class ClientTimeout:
    total: Optional[float] = None
    connect: Optional[float] = None
    sock_read: Optional[float] = None
    sock_connect: Optional[float] = None

    # pool_queue_timeout: Optional[float] = None
    # dns_resolution_timeout: Optional[float] = None
    # socket_connect_timeout: Optional[float] = None
    # connection_acquiring_timeout: Optional[float] = None
    # new_connection_timeout: Optional[float] = None
    # http_header_timeout: Optional[float] = None
    # response_body_timeout: Optional[float] = None

    # to create a timeout specific for a single request, either
    # - create a completely new one to overwrite the default
    # - or use http://www.attrs.org/en/stable/api.html#attr.evolve
    # to overwrite the defaults


# 5 Minute default read timeout
DEFAULT_TIMEOUT: Final[ClientTimeout] = ClientTimeout(total=5 * 60)

_RetType = TypeVar("_RetType")


class ClientSession:
    """First-class interface for making HTTP requests."""

    ATTRS = frozenset(
        [
            "_base_url",
            "_source_traceback",
            "_connector",
            "requote_redirect_url",
            "_loop",
            "_cookie_jar",
            "_connector_owner",
            "_default_auth",
            "_version",
            "_json_serialize",
            "_requote_redirect_url",
            "_timeout",
            "_raise_for_status",
            "_auto_decompress",
            "_trust_env",
            "_default_headers",
            "_skip_auto_headers",
            "_request_class",
            "_response_class",
            "_ws_response_class",
            "_trace_configs",
            "_read_bufsize",
        ]
    )

    _source_traceback = None  # type: Optional[traceback.StackSummary]
    _connector = None  # type: Optional[BaseConnector]

    def __init__(
        self,
        base_url: Optional[StrOrURL] = None,
        *,
        connector: Optional[BaseConnector] = None,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        cookies: Optional[LooseCookies] = None,
        headers: Optional[LooseHeaders] = None,
        skip_auto_headers: Optional[Iterable[str]] = None,
        auth: Optional[BasicAuth] = None,
        json_serialize: JSONEncoder = json.dumps,
        request_class: Type[ClientRequest] = ClientRequest,
        response_class: Type[ClientResponse] = ClientResponse,
        ws_response_class: Type[ClientWebSocketResponse] = ClientWebSocketResponse,
        version: HttpVersion = http.HttpVersion11,
        cookie_jar: Optional[AbstractCookieJar] = None,
        connector_owner: bool = True,
        raise_for_status: bool = False,
        read_timeout: Union[float, object] = sentinel,
        conn_timeout: Optional[float] = None,
        timeout: Union[object, ClientTimeout] = sentinel,
        auto_decompress: bool = True,
        trust_env: bool = False,
        requote_redirect_url: bool = True,
        trace_configs: Optional[List[TraceConfig]] = None,
        read_bufsize: int = 2**16,
    ) -> None:
        if loop is None:
            if connector is not None:
                loop = connector._loop

        loop = get_running_loop(loop)

        if base_url is None or isinstance(base_url, URL):
            self._base_url: Optional[URL] = base_url
        else:
            self._base_url = URL(base_url)
            assert (
                self._base_url.origin() == self._base_url
            ), "Only absolute URLs without path part are supported"

        if connector is None:
            connector = TCPConnector(loop=loop)

        if connector._loop is not loop:
            raise RuntimeError("Session and connector has to use same event loop")

        self._loop = loop

        if loop.get_debug():
            self._source_traceback = traceback.extract_stack(sys._getframe(1))

        if cookie_jar is None:
            cookie_jar = CookieJar(loop=loop)
        self._cookie_jar = cookie_jar

        if cookies is not None:
            self._cookie_jar.update_cookies(cookies)

        self._connector = connector
        self._connector_owner = connector_owner
        self._default_auth = auth
        self._version = version
        self._json_serialize = json_serialize
        if timeout is sentinel:
            self._timeout = DEFAULT_TIMEOUT
            if read_timeout is not sentinel:
                warnings.warn(
                    "read_timeout is deprecated, " "use timeout argument instead",
                    DeprecationWarning,
                    stacklevel=2,
                )
                self._timeout = attr.evolve(self._timeout, total=read_timeout)
            if conn_timeout is not None:
                self._timeout = attr.evolve(self._timeout, connect=conn_timeout)
                warnings.warn(
                    "conn_timeout is deprecated, " "use timeout argument instead",
                    DeprecationWarning,
                    stacklevel=2,
                )
        else:
            self._timeout = timeout  # type: ignore[assignment]
            if read_timeout is not sentinel:
                raise ValueError(
                    "read_timeout and timeout parameters "
                    "conflict, please setup "
                    "timeout.read"
                )
            if conn_timeout is not None:
                raise ValueError(
                    "conn_timeout and timeout parameters "
                    "conflict, please setup "
                    "timeout.connect"
                )
        self._raise_for_status = raise_for_status
        self._auto_decompress = auto_decompress
        self._trust_env = trust_env
        self._requote_redirect_url = requote_redirect_url
        self._read_bufsize = read_bufsize

        # Convert to list of tuples
        if headers:
            real_headers: CIMultiDict[str] = CIMultiDict(headers)
        else:
            real_headers = CIMultiDict()
        self._default_headers: CIMultiDict[str] = real_headers
        if skip_auto_headers is not None:
            self._skip_auto_headers = frozenset(istr(i) for i in skip_auto_headers)
        else:
            self._skip_auto_headers = frozenset()

        self._request_class = request_class
        self._response_class = response_class
        self._ws_response_class = ws_response_class

        self._trace_configs = trace_configs or []
        for trace_config in self._trace_configs:
            trace_config.freeze()

    def __init_subclass__(cls: Type["ClientSession"]) -> None:
        warnings.warn(
            "Inheritance class {} from ClientSession "
            "is discouraged".format(cls.__name__),
            DeprecationWarning,
            stacklevel=2,
        )

    if DEBUG:

        def __setattr__(self, name: str, val: Any) -> None:
            if name not in self.ATTRS:
                warnings.warn(
                    "Setting custom ClientSession.{} attribute "
                    "is discouraged".format(name),
                    DeprecationWarning,
                    stacklevel=2,
                )
            super().__setattr__(name, val)

    def __del__(self, _warnings: Any = warnings) -> None:
        if not self.closed:
            if PY_36:
                kwargs = {"source": self}
            else:
                kwargs = {}
            _warnings.warn(
                f"Unclosed client session {self!r}", ResourceWarning, **kwargs
            )
            context = {"client_session": self, "message": "Unclosed client session"}
            if self._source_traceback is not None:
                context["source_traceback"] = self._source_traceback
            self._loop.call_exception_handler(context)

    def request(
        self, method: str, url: StrOrURL, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP request."""
        return _RequestContextManager(self._request(method, url, **kwargs))

    def _build_url(self, str_or_url: StrOrURL) -> URL:
        url = URL(str_or_url)
        if self._base_url is None:
            return url
        else:
            assert not url.is_absolute() and url.path.startswith("/")
            return self._base_url.join(url)

    async def _request(
        self,
        method: str,
        str_or_url: StrOrURL,
        *,
        name: str = None,
        params: Optional[Mapping[str, str]] = None,
        data: Any = None,
        json: Any = None,
        cookies: Optional[LooseCookies] = None,
        headers: Optional[LooseHeaders] = None,
        skip_auto_headers: Optional[Iterable[str]] = None,
        auth: Optional[BasicAuth] = None,
        allow_redirects: bool = True,
        max_redirects: int = 10,
        compress: Optional[str] = None,
        chunked: Optional[bool] = None,
        expect100: bool = False,
        raise_for_status: Optional[bool] = None,
        read_until_eof: bool = True,
        proxy: Optional[StrOrURL] = None,
        proxy_auth: Optional[BasicAuth] = None,
        timeout: Union[ClientTimeout, object] = sentinel,
        verify_ssl: Optional[bool] = None,
        fingerprint: Optional[bytes] = None,
        ssl_context: Optional[SSLContext] = None,
        ssl: Optional[Union[SSLContext, bool, Fingerprint]] = None,
        proxy_headers: Optional[LooseHeaders] = None,
        trace_request_ctx: Optional[SimpleNamespace] = None,
        read_bufsize: Optional[int] = None,
    ) -> ClientResponse:

        # NOTE: timeout clamps existing connect and read timeouts.  We cannot
        # set the default to None because we need to detect if the user wants
        # to use the existing timeouts by setting timeout to None.

        if self.closed:
            raise RuntimeError("Session is closed")

        ssl = _merge_ssl_params(ssl, verify_ssl, ssl_context, fingerprint)

        if data is not None and json is not None:
            raise ValueError(
                "data and json parameters can not be used at the same time"
            )
        elif json is not None:
            data = payload.JsonPayload(json, dumps=self._json_serialize)

        if not isinstance(chunked, bool) and chunked is not None:
            warnings.warn("Chunk size is deprecated #1615", DeprecationWarning)

        redirects = 0
        history = []
        version = self._version

        # Merge with default headers and transform to CIMultiDict
        headers = self._prepare_headers(headers)
        proxy_headers = self._prepare_headers(proxy_headers)

        try:
            url = self._build_url(str_or_url)
        except ValueError as e:
            raise InvalidURL(str_or_url) from e

        skip_headers = set(self._skip_auto_headers)
        if skip_auto_headers is not None:
            for i in skip_auto_headers:
                skip_headers.add(istr(i))

        if proxy is not None:
            try:
                proxy = URL(proxy)
            except ValueError as e:
                raise InvalidURL(proxy) from e

        if timeout is sentinel:
            real_timeout: ClientTimeout = self._timeout
        else:
            if not isinstance(timeout, ClientTimeout):
                real_timeout = ClientTimeout(total=timeout)  # type: ignore[arg-type]
            else:
                real_timeout = timeout
        # timeout is cumulative for all request operations
        # (request, redirects, responses, data consuming)
        tm = TimeoutHandle(self._loop, real_timeout.total)
        handle = tm.start()

        if read_bufsize is None:
            read_bufsize = self._read_bufsize

        traces = [
            Trace(
                self,
                trace_config,
                trace_config.trace_config_ctx(trace_request_ctx=trace_request_ctx),
            )
            for trace_config in self._trace_configs
        ]

        for trace in traces:
            await trace.send_request_start(method, url.update_query(params), headers)

        timer = tm.timer()
        try:
            with timer:
                while True:
                    url, auth_from_url = strip_auth_from_url(url)
                    if auth and auth_from_url:
                        raise ValueError(
                            "Cannot combine AUTH argument with "
                            "credentials encoded in URL"
                        )

                    if auth is None:
                        auth = auth_from_url
                    if auth is None:
                        auth = self._default_auth
                    # It would be confusing if we support explicit
                    # Authorization header with auth argument
                    if (
                        headers is not None
                        and auth is not None
                        and hdrs.AUTHORIZATION in headers
                    ):
                        raise ValueError(
                            "Cannot combine AUTHORIZATION header "
                            "with AUTH argument or credentials "
                            "encoded in URL"
                        )

                    all_cookies = self._cookie_jar.filter_cookies(url)

                    if cookies is not None:
                        tmp_cookie_jar = CookieJar()
                        tmp_cookie_jar.update_cookies(cookies)
                        req_cookies = tmp_cookie_jar.filter_cookies(url)
                        if req_cookies:
                            all_cookies.load(req_cookies)

                    if proxy is not None:
                        proxy = URL(proxy)
                    elif self._trust_env:
                        with suppress(LookupError):
                            proxy, proxy_auth = get_env_proxy_for_url(url)

                    req = self._request_class(
                        method,
                        url,
                        params=params,
                        headers=headers,
                        skip_auto_headers=skip_headers,
                        data=data,
                        cookies=all_cookies,
                        auth=auth,
                        version=version,
                        compress=compress,
                        chunked=chunked,
                        expect100=expect100,
                        loop=self._loop,
                        response_class=self._response_class,
                        proxy=proxy,
                        proxy_auth=proxy_auth,
                        timer=timer,
                        session=self,
                        ssl=ssl,
                        proxy_headers=proxy_headers,
                        traces=traces,
                    )
                    # Collect request data
                    request_meta = {}
                    request_meta["request_name"] = name if name else url.path
                    request_meta["request_method"] = method
                    
                    # connection timeout
                    try:
                        async with ceil_timeout(real_timeout.connect):
                            assert self._connector is not None
                            # requests start time
                            request_meta["start_time"] = default_timer()
                            conn = await self._connector.connect(
                                req, traces=traces, timeout=real_timeout
                            )
                    except asyncio.TimeoutError as exc:
                        raise ServerTimeoutError(
                            "Connection timeout " "to host {}".format(url)
                        ) from exc

                    assert conn.transport is not None

                    assert conn.protocol is not None
                    conn.protocol.set_response_params(
                        timer=timer,
                        skip_payload=method.upper() == "HEAD",
                        read_until_eof=read_until_eof,
                        auto_decompress=self._auto_decompress,
                        read_timeout=real_timeout.sock_read,
                        read_bufsize=read_bufsize,
                    )

                    try:
                        try:
                            resp = await req.send(conn)
                            try:
                                await resp.start(conn)
                            except BaseException:
                                resp.close()
                                raise
                        except BaseException:
                            conn.close()
                            raise
                    except ClientError:
                        raise
                    except OSError as exc:
                        if exc.errno is None and isinstance(exc, asyncio.TimeoutError):
                            raise
                        raise ClientOSError(*exc.args) from exc

                    self._cookie_jar.update_cookies(resp.cookies, resp.url)

                    # redirects
                    if resp.status in (301, 302, 303, 307, 308) and allow_redirects:

                        for trace in traces:
                            await trace.send_request_redirect(
                                method, url.update_query(params), headers, resp
                            )

                        redirects += 1
                        history.append(resp)
                        if max_redirects and redirects >= max_redirects:
                            resp.close()
                            raise TooManyRedirects(
                                history[0].request_info, tuple(history)
                            )

                        # For 301 and 302, mimic IE, now changed in RFC
                        # https://github.com/kennethreitz/requests/pull/269
                        if (resp.status == 303 and resp.method != hdrs.METH_HEAD) or (
                            resp.status in (301, 302) and resp.method == hdrs.METH_POST
                        ):
                            method = hdrs.METH_GET
                            data = None
                            if headers.get(hdrs.CONTENT_LENGTH):
                                headers.pop(hdrs.CONTENT_LENGTH)

                        r_url = resp.headers.get(hdrs.LOCATION) or resp.headers.get(
                            hdrs.URI
                        )
                        if r_url is None:
                            # see github.com/aio-libs/aiohttp/issues/2022
                            break
                        else:
                            # reading from correct redirection
                            # response is forbidden
                            resp.release()

                        try:
                            parsed_url = URL(
                                r_url, encoded=not self._requote_redirect_url
                            )

                        except ValueError as e:
                            raise InvalidURL(r_url) from e

                        scheme = parsed_url.scheme
                        if scheme not in ("http", "https", ""):
                            resp.close()
                            raise ValueError("Can redirect only to http or https")
                        elif not scheme:
                            parsed_url = url.join(parsed_url)

                        if url.origin() != parsed_url.origin():
                            auth = None
                            headers.pop(hdrs.AUTHORIZATION, None)

                        url = parsed_url
                        params = None
                        resp.release()
                        continue

                    break

            # check response status
            if raise_for_status is None:
                raise_for_status = self._raise_for_status
            if raise_for_status:
                resp.raise_for_status()

            # register connection
            if handle is not None:
                if resp.connection is not None:
                    resp.connection.add_callback(handle.cancel)
                else:
                    handle.cancel()

            resp._history = tuple(history)
            

            for trace in traces:
                await trace.send_request_end(
                    method, url.update_query(params), headers, resp
                )
            resp.request_meta = request_meta
            return resp

        except BaseException as e:
            # cleanup timer
            tm.close()
            if handle:
                handle.cancel()
                handle = None

            for trace in traces:
                await trace.send_request_exception(
                    method, url.update_query(params), headers, e
                )
            raise

    def ws_connect(
        self,
        url: StrOrURL,
        *,
        method: str = hdrs.METH_GET,
        protocols: Iterable[str] = (),
        timeout: float = 10.0,
        receive_timeout: Optional[float] = None,
        autoclose: bool = True,
        autoping: bool = True,
        heartbeat: Optional[float] = None,
        auth: Optional[BasicAuth] = None,
        origin: Optional[str] = None,
        params: Optional[Mapping[str, str]] = None,
        headers: Optional[LooseHeaders] = None,
        proxy: Optional[StrOrURL] = None,
        proxy_auth: Optional[BasicAuth] = None,
        ssl: Union[SSLContext, bool, None, Fingerprint] = None,
        verify_ssl: Optional[bool] = None,
        fingerprint: Optional[bytes] = None,
        ssl_context: Optional[SSLContext] = None,
        proxy_headers: Optional[LooseHeaders] = None,
        compress: int = 0,
        max_msg_size: int = 4 * 1024 * 1024,
    ) -> "_WSRequestContextManager":
        """Initiate websocket connection."""
        return _WSRequestContextManager(
            self._ws_connect(
                url,
                method=method,
                protocols=protocols,
                timeout=timeout,
                receive_timeout=receive_timeout,
                autoclose=autoclose,
                autoping=autoping,
                heartbeat=heartbeat,
                auth=auth,
                origin=origin,
                params=params,
                headers=headers,
                proxy=proxy,
                proxy_auth=proxy_auth,
                ssl=ssl,
                verify_ssl=verify_ssl,
                fingerprint=fingerprint,
                ssl_context=ssl_context,
                proxy_headers=proxy_headers,
                compress=compress,
                max_msg_size=max_msg_size,
            )
        )

    async def _ws_connect(
        self,
        url: StrOrURL,
        *,
        method: str = hdrs.METH_GET,
        protocols: Iterable[str] = (),
        timeout: float = 10.0,
        receive_timeout: Optional[float] = None,
        autoclose: bool = True,
        autoping: bool = True,
        heartbeat: Optional[float] = None,
        auth: Optional[BasicAuth] = None,
        origin: Optional[str] = None,
        params: Optional[Mapping[str, str]] = None,
        headers: Optional[LooseHeaders] = None,
        proxy: Optional[StrOrURL] = None,
        proxy_auth: Optional[BasicAuth] = None,
        ssl: Union[SSLContext, bool, None, Fingerprint] = None,
        verify_ssl: Optional[bool] = None,
        fingerprint: Optional[bytes] = None,
        ssl_context: Optional[SSLContext] = None,
        proxy_headers: Optional[LooseHeaders] = None,
        compress: int = 0,
        max_msg_size: int = 4 * 1024 * 1024,
    ) -> ClientWebSocketResponse:

        if headers is None:
            real_headers: CIMultiDict[str] = CIMultiDict()
        else:
            real_headers = CIMultiDict(headers)

        default_headers = {
            hdrs.UPGRADE: "websocket",
            hdrs.CONNECTION: "upgrade",
            hdrs.SEC_WEBSOCKET_VERSION: "13",
        }

        for key, value in default_headers.items():
            real_headers.setdefault(key, value)

        sec_key = base64.b64encode(os.urandom(16))
        real_headers[hdrs.SEC_WEBSOCKET_KEY] = sec_key.decode()

        if protocols:
            real_headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ",".join(protocols)
        if origin is not None:
            real_headers[hdrs.ORIGIN] = origin
        if compress:
            extstr = ws_ext_gen(compress=compress)
            real_headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = extstr

        ssl = _merge_ssl_params(ssl, verify_ssl, ssl_context, fingerprint)

        # send request
        resp = await self.request(
            method,
            url,
            params=params,
            headers=real_headers,
            read_until_eof=False,
            auth=auth,
            proxy=proxy,
            proxy_auth=proxy_auth,
            ssl=ssl,
            proxy_headers=proxy_headers,
        )

        try:
            # check handshake
            if resp.status != 101:
                raise WSServerHandshakeError(
                    resp.request_info,
                    resp.history,
                    message="Invalid response status",
                    status=resp.status,
                    headers=resp.headers,
                )

            if resp.headers.get(hdrs.UPGRADE, "").lower() != "websocket":
                raise WSServerHandshakeError(
                    resp.request_info,
                    resp.history,
                    message="Invalid upgrade header",
                    status=resp.status,
                    headers=resp.headers,
                )

            if resp.headers.get(hdrs.CONNECTION, "").lower() != "upgrade":
                raise WSServerHandshakeError(
                    resp.request_info,
                    resp.history,
                    message="Invalid connection header",
                    status=resp.status,
                    headers=resp.headers,
                )

            # key calculation
            r_key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, "")
            match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()).decode()
            if r_key != match:
                raise WSServerHandshakeError(
                    resp.request_info,
                    resp.history,
                    message="Invalid challenge response",
                    status=resp.status,
                    headers=resp.headers,
                )

            # websocket protocol
            protocol = None
            if protocols and hdrs.SEC_WEBSOCKET_PROTOCOL in resp.headers:
                resp_protocols = [
                    proto.strip()
                    for proto in resp.headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(",")
                ]

                for proto in resp_protocols:
                    if proto in protocols:
                        protocol = proto
                        break

            # websocket compress
            notakeover = False
            if compress:
                compress_hdrs = resp.headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS)
                if compress_hdrs:
                    try:
                        compress, notakeover = ws_ext_parse(compress_hdrs)
                    except WSHandshakeError as exc:
                        raise WSServerHandshakeError(
                            resp.request_info,
                            resp.history,
                            message=exc.args[0],
                            status=resp.status,
                            headers=resp.headers,
                        ) from exc
                else:
                    compress = 0
                    notakeover = False

            conn = resp.connection
            assert conn is not None
            conn_proto = conn.protocol
            assert conn_proto is not None
            transport = conn.transport
            assert transport is not None
            reader: FlowControlDataQueue[WSMessage] = FlowControlDataQueue(
                conn_proto, 2**16, loop=self._loop
            )
            conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader)
            writer = WebSocketWriter(
                conn_proto,
                transport,
                use_mask=True,
                compress=compress,
                notakeover=notakeover,
            )
        except BaseException:
            resp.close()
            raise
        else:
            return self._ws_response_class(
                reader,
                writer,
                protocol,
                resp,
                timeout,
                autoclose,
                autoping,
                self._loop,
                receive_timeout=receive_timeout,
                heartbeat=heartbeat,
                compress=compress,
                client_notakeover=notakeover,
            )

    def _prepare_headers(self, headers: Optional[LooseHeaders]) -> "CIMultiDict[str]":
        """Add default headers and transform it to CIMultiDict"""
        # Convert headers to MultiDict
        result = CIMultiDict(self._default_headers)
        if headers:
            if not isinstance(headers, (MultiDictProxy, MultiDict)):
                headers = CIMultiDict(headers)
            added_names: Set[str] = set()
            for key, value in headers.items():
                if key in added_names:
                    result.add(key, value)
                else:
                    result[key] = value
                    added_names.add(key)
        return result

    def get(
        self, url: StrOrURL, *, allow_redirects: bool = True, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP GET request."""
        return _RequestContextManager(
            self._request(hdrs.METH_GET, url, allow_redirects=allow_redirects, **kwargs)
        )

    def options(
        self, url: StrOrURL, *, allow_redirects: bool = True, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP OPTIONS request."""
        return _RequestContextManager(
            self._request(
                hdrs.METH_OPTIONS, url, allow_redirects=allow_redirects, **kwargs
            )
        )

    def head(
        self, url: StrOrURL, *, allow_redirects: bool = False, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP HEAD request."""
        return _RequestContextManager(
            self._request(
                hdrs.METH_HEAD, url, allow_redirects=allow_redirects, **kwargs
            )
        )

    def post(
        self, url: StrOrURL, *, data: Any = None, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP POST request."""
        return _RequestContextManager(
            self._request(hdrs.METH_POST, url, data=data, **kwargs)
        )

    def put(
        self, url: StrOrURL, *, data: Any = None, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP PUT request."""
        return _RequestContextManager(
            self._request(hdrs.METH_PUT, url, data=data, **kwargs)
        )

    def patch(
        self, url: StrOrURL, *, data: Any = None, **kwargs: Any
    ) -> "_RequestContextManager":
        """Perform HTTP PATCH request."""
        return _RequestContextManager(
            self._request(hdrs.METH_PATCH, url, data=data, **kwargs)
        )

    def delete(self, url: StrOrURL, **kwargs: Any) -> "_RequestContextManager":
        """Perform HTTP DELETE request."""
        return _RequestContextManager(self._request(hdrs.METH_DELETE, url, **kwargs))

    async def close(self) -> None:
        """Close underlying connector.

        Release all acquired resources.
        """
        if not self.closed:
            if self._connector is not None and self._connector_owner:
                await self._connector.close()
            self._connector = None

    @property
    def closed(self) -> bool:
        """Is client session closed.

        A readonly property.
        """
        return self._connector is None or self._connector.closed

    @property
    def connector(self) -> Optional[BaseConnector]:
        """Connector instance used for the session."""
        return self._connector

    @property
    def cookie_jar(self) -> AbstractCookieJar:
        """The session cookies."""
        return self._cookie_jar

    @property
    def version(self) -> Tuple[int, int]:
        """The session HTTP protocol version."""
        return self._version

    @property
    def requote_redirect_url(self) -> bool:
        """Do URL requoting on redirection handling."""
        return self._requote_redirect_url

    @requote_redirect_url.setter
    def requote_redirect_url(self, val: bool) -> None:
        """Do URL requoting on redirection handling."""
        warnings.warn(
            "session.requote_redirect_url modification " "is deprecated #2778",
            DeprecationWarning,
            stacklevel=2,
        )
        self._requote_redirect_url = val

    @property
    def loop(self) -> asyncio.AbstractEventLoop:
        """Session's loop."""
        warnings.warn(
            "client.loop property is deprecated", DeprecationWarning, stacklevel=2
        )
        return self._loop

    @property
    def timeout(self) -> ClientTimeout:
        """Timeout for the session."""
        return self._timeout

    @property
    def headers(self) -> "CIMultiDict[str]":
        """The default headers of the client session."""
        return self._default_headers

    @property
    def skip_auto_headers(self) -> FrozenSet[istr]:
        """Headers for which autogeneration should be skipped"""
        return self._skip_auto_headers

    @property
    def auth(self) -> Optional[BasicAuth]:
        """An object that represents HTTP Basic Authorization"""
        return self._default_auth

    @property
    def json_serialize(self) -> JSONEncoder:
        """Json serializer callable"""
        return self._json_serialize

    @property
    def connector_owner(self) -> bool:
        """Should connector be closed on session closing"""
        return self._connector_owner

    @property
    def raise_for_status(
        self,
    ) -> Union[bool, Callable[[ClientResponse], Awaitable[None]]]:
        """Should `ClientResponse.raise_for_status()` be called for each response."""
        return self._raise_for_status

    @property
    def auto_decompress(self) -> bool:
        """Should the body response be automatically decompressed."""
        return self._auto_decompress

    @property
    def trust_env(self) -> bool:
        """
        Should proxies information from environment or netrc be trusted.

        Information is from HTTP_PROXY / HTTPS_PROXY environment variables
        or ~/.netrc file if present.
        """
        return self._trust_env

    @property
    def trace_configs(self) -> List[TraceConfig]:
        """A list of TraceConfig instances used for client tracing"""
        return self._trace_configs

    def detach(self) -> None:
        """Detach connector from session without closing the former.

        Session is switched to closed state anyway.
        """
        self._connector = None

    def __enter__(self) -> None:
        raise TypeError("Use async with instead")

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        # __exit__ should exist in pair with __enter__ but never executed
        pass  # pragma: no cover

    async def __aenter__(self) -> "ClientSession":
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        await self.close()


class _BaseRequestContextManager(Coroutine[Any, Any, _RetType], Generic[_RetType]):

    __slots__ = ("_coro", "_resp", "_manual_result")

    def __init__(self, coro: Coroutine["asyncio.Future[Any]", None, _RetType]) -> None:
        self._coro = coro
        self._manual_result = None

    def send(self, arg: None) -> "asyncio.Future[Any]":
        return self._coro.send(arg)

    def throw(self, arg: BaseException) -> None:  # type: ignore[arg-type,override]
        self._coro.throw(arg)

    def close(self) -> None:
        return self._coro.close()

    def __await__(self) -> Generator[Any, None, _RetType]:
        ret = self._coro.__await__()
        return ret

    def __iter__(self) -> Generator[Any, None, _RetType]:
        return self.__await__()

    async def __aenter__(self) -> _RetType:
        self._resp = await self._coro
        self._resp.success = self.success
        self._resp.failure = self.failure
        # requests meta data
        if self._resp.headers.get("Transfer-Encoding") != "chunked":
            self._resp.request_meta["response_time"] = int((default_timer() - self._resp.request_meta["start_time"]) * 1000)
            self._resp.request_meta["response_length"] = float(self._resp.headers.get("Content-Length"))
            self._resp.request_meta.pop("start_time")
        return self._resp

    def success(self):
        """
        Report the response as successful
        
        Example::
        
            async with self.sessen.get("/does/not/exist") as resp:
                if resp.status == 404:
                    resp.success()
        """
        self._manual_result = True
    
    def failure(self, exc):
        """
        Report the response as a failure.
        
        exc can be either a python exception, or a string in which case it will
        be wrapped inside a CatchResponseError. 
        
        Example::
        
            async with self.sessen.get("/") as resp:
                if resp.status == 200:
                    response.failure("is 200")
        """
        if not isinstance(exc, Exception):
            exc = CatchResponseError(exc)
        self._manual_result = exc


class _RequestContextManager(_BaseRequestContextManager[ClientResponse]):
    __slots__ = ()
   
    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        # We're basing behavior on the exception as it can be caused by
        # user code unrelated to the status of the connection.  If you
        # would like to close a connection you must do that
        # explicitly.  Otherwise connection error handling should kick in
        # and close/recycle the connection as required.

        # requests meta data
        if self._resp.request_meta.get("start_time"):
            self._resp.request_meta["response_time"] = int((default_timer() - self._resp.request_meta["start_time"]) * 1000)
            self._resp.request_meta["response_length"] = len(self._resp._body or "")
            self._resp.request_meta.pop("start_time")
        self._resp.request_meta["error"] = None
        
        if self._manual_result:
            if self._manual_result is True:
                await events.stats_request.fire(runner=runners.global_runner, **self._resp.request_meta)
            elif isinstance(self._manual_result, Exception):
                self._resp.request_meta["error"] = self._manual_result
                await events.stats_request.fire(runner=runners.global_runner, **self._resp.request_meta)
            self._resp.release()
            # we want other unknown exceptions to be raised
            return exc_type is None
        if exc_type:
            if isinstance(exc, REQUESTERRORS):
                self._resp.request_meta["error"] = f"{sys.exc_info()[0].__name__}: {exc}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip()
                await events.stats_request.fire(runner=runners.global_runner, **self._resp.request_meta)
            else:
                self._resp.release()
                # we want other unknown exceptions to be raised
                return False
        else:
            try:
                self._resp.raise_for_status()
            except ClientResponseError as e:
                self._resp.request_meta["error"] = e
            await events.stats_request.fire(runner=runners.global_runner, **self._resp.request_meta)
        
        self._resp.release()
        return True


class _WSRequestContextManager(_BaseRequestContextManager[ClientWebSocketResponse]):
    __slots__ = ()

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        await self._resp.close()




aiotest/client.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import traceback
from aiotest.log import logger


class EventHook:
    
    def __init__(self):
        self._handlers = []
    
    def __iadd__(self, handler):
        self._handlers.append(handler)
        return self
    
    def __isub__(self, handler):
        self._handlers.remove(handler)
        return self
    
    async def fire(self, **kwargs):

        for handler in self._handlers:
            try:
                await handler(**kwargs)
            except Exception as e:
                logger.error(f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

init_command_line_parser = EventHook()
init = EventHook()
test_start = EventHook()
start_complete = EventHook()
stats_request = EventHook()
user_error = EventHook()
worker_report = EventHook()
test_stop = EventHook()
quitting = EventHook()


aiotest/events.py
hewei198711/Aiotest
# encoding: utf-8


class RPCError(Exception):
    pass


class CatchResponseError(Exception):
    pass
aiotest/exception.py
hewei198711/Aiotest
# encoding: utf-8

import sys

from loguru import logger


def setup_logging(loglevel, logfile=None):
    loglevel = loglevel.upper()
    logger.remove()
    logger.add(sys.stderr, level=loglevel)
    if logfile:
        logger.add(logfile, level=loglevel, retention="1 days")
aiotest/log.py
hewei198711/Aiotest
# encoding: utf-8

import asyncio
import re
import os
import sys
import inspect
import argparse
import importlib
from datetime import timedelta
from pprint import pprint

import aiotest
from prometheus_client import Histogram
from aiotest import stats_exporter
from aiotest.stats_exporter import STATSERROR
from aiotest.log import logger, setup_logging
from aiotest.shape import LoadUserShape
from aiotest.user import User
from aiotest import events
from aiotest import runners
from aiotest.runners import LocalRunner, MasterRunner, WorkerRunner


version = aiotest.__version__

async def parse_options(args=None):
    """
    Handle command-line options with argparse.ArgumentParser.

    Return list of arguments, largely for use in `parse_arguments`.
    """

    # Initialize
    parser = argparse.ArgumentParser(
        prog="aiotest [options]", 
        description="asyncio is an easy to use, scriptable and scalable performance testing tool"
    )

    parser.add_argument(
        '-f', '--aiotestfile',
        default='aiotestfile',
        help="Python module file to import, e.g. '../other.py'. Default: aiotestfile"
    )
        
    parser.add_argument(
        '-H', '--host',
        default= "",
        help="Host to load test in the following format: http://10.21.32.33"
    )
    
    # master-worker
    parser.add_argument(
        '--master',
        action='store_true',
        help="Set aiotest to run in distributed mode with this process as master"
    )

    parser.add_argument(
        '--worker',
        action='store_true',
        help="Set aiotest to run in distributed mode with this process as worker"
    )

    parser.add_argument(
        '--expect-workers',
        type=int,
        default=1,
        help="How many workers master should expect to connect before starting the test"
    )
        
    parser.add_argument(
        '--master-host',
        default="127.0.0.1",
        help="Host or IP address of aiotest master for distributed load testing. Only used when running with --worker. Defaults to 127.0.0.1."
    )
    
    parser.add_argument(
        '--master-port',
        type=int,
        default=5557,
        help="The port to connect to that is used by the aiotest master for distributed load testing. Only used when running with --worker. Defaults to 5557. Note that workers will also connect to the master node on this port + 1."
    )

    parser.add_argument(
        '--master-bind-host',
        default="*",
        help="Interfaces (hostname, ip) that aiotest master should bind to. Only used when running with --master. Defaults to * (all available interfaces)."
    )
    
    parser.add_argument(
        '--master-bind-port',
        type=int,
        default=5557,
        help="Port that worker master should bind to. Only used when running with --master. Defaults to 5557. Note that aiotest will also use this port + 1, so by default the master node will bind to 5557 and 5558."
    )

    # users rate run_time
    parser.add_argument(
        '-u', '--users',
        type=int,
        dest='num_users',
        default=1,
        help="Number of concurrent users"
    )

    parser.add_argument(
        '-r', '--rate',
        type=int,
        default=1,
        help="The rate per second in which users are started/stoped"
    )

    parser.add_argument(
        '-t', '--run-time',
        type=str,
        default=None,
        help="Stop after the specified amount of time, e.g. (300s, 20m, 3h, 1h30m, etc.)"
    )    

    # prometheus
    parser.add_argument(
        '--prometheus-port',
        type=int,
        default=8089,
        help="Port that metrics are exposed over HTTP, to be read by the Prometheus server."
    )
    
    parser.add_argument(
        '-b', '--buckets', 
        type=int, nargs='*', 
        default= [50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 2000, 5000, 8000], 
        help='Prometheus histogram buckets'
    )    
    
    # loglevel logfile
    parser.add_argument(
        '--loglevel', '-L',
        default='INFO',
        help="Choose between DEBUG/INFO/WARNING/ERROR/CRITICAL. Default is INFO.",
    )
    
    parser.add_argument(
        '--logfile',
        default = None,
        help="Path to log file. If not set, log will go to stdout/stderr",
    )
    
    # show 
    parser.add_argument(
        '--show-users-wight',
        action='store_true',
        help="print json data of the users classes' execution wight"
    )
    
    parser.add_argument(
        '-V', '--version',
        action='version',
        version= f'aiotest=={version}',
    )

    # Provides a hook to add command-line arguments
    await events.init_command_line_parser.fire(parser=parser)
   
    return parser.parse_args(args=args)

def find_aiotestfile(aiotestfile):
    """
    Attempt to locate a aiotestfile by either explicitly.
    """
    _, suffix = os.path.splitext(aiotestfile)
    if suffix and suffix != ".py":
        logger.warning(f"{aiotestfile} mast be .py file")
        sys.exit(1)
    if not suffix:
        aiotestfile += ".py"
    # return abspath
    if os.path.dirname(aiotestfile):
        # expand home-directory markers and test for existence
        expanded = os.path.expanduser(aiotestfile)
        if os.path.exists(expanded):
            return os.path.abspath(expanded)
    else:
        path = os.path.abspath(".")
        joined = os.path.join(path, aiotestfile)
        if os.path.exists(joined):
            return os.path.abspath(joined)
    
    logger.warning(f"Current working directory:{os.path.abspath('.')}，File not found:{aiotestfile}")

def is_shape_class(cls):
    "Check if a class is a SubLoadUserShape"
    return bool(
        inspect.isclass(cls)
        and issubclass(cls, LoadUserShape)
        and (cls.__name__.startswith("Test") or cls.__name__.endswith("Test"))
        and cls.__dict__["__module__"] != "aiotest.shape"
    )

def is_user_class(cls):
    """
    Check if a variable is a runnable (non-abstract) SubUser class
    """
    return bool(
        inspect.isclass(cls)
        and issubclass(cls, User)
        and getattr(cls, "weight")
        and (cls.__name__.startswith("Test") or cls.__name__.endswith("Test"))
        and cls.__dict__["__module__"] != "aiotest.user"
        )

def load_aiotestfile(path):
    """
    Import given aiotestfile path and return (UserClass, LoadUserShapeClass).
    """
    # Start with making sure the current working dir is in the sys.path
    sys.path.insert(0, os.getcwd())
    directory, aiotestfile = os.path.split(path)
    # If the directory isn't in the PYTHONPATH, add it so our import will work
    added_to_path = False
    index = None
    if directory not in sys.path:
        sys.path.insert(0, directory)
        added_to_path = True
    # If the directory IS in the PYTHONPATH, move it to the front temporarily,
    # otherwise other aiotestfiles -- like Aiotest's own -- may scoop the intended one.
    else:
        i = sys.path.index(directory)
        if i != 0:
            index = i
            # Add to front, then remove from original position
            sys.path.insert(0, directory)
            del sys.path[i+1]

    imported = importlib.import_module(os.path.splitext(aiotestfile)[0])
    # Remove directory from path if we added it ourselves (just to be neat)
    if added_to_path:
        del sys.path[0]
    # Put back in original index if we moved it
    if index:
        del sys.path[0]
        sys.path.insert(index, directory)
    # Return our two-tuple
    user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)}
    # Find shape class, if any, return it
    shape_classes = [value for value in vars(imported).values() if is_shape_class(value)]
    if len(shape_classes) == 1:
        shape_class = shape_classes[0]()
    elif len(shape_classes) > 1:
        logger.error("tactics class  only have one!!")
        sys.exit(1)
    else:
        shape_class = None
    return user_classes, shape_class

def parse_timespan(time_str):
    """
    Parse a string representing a time span and return the number of seconds.
    Valid formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.
    """
    if not time_str:
        raise ValueError("Invalid time span format")
    
    if re.match(r"^\d+$", time_str):
        # if an int is specified we assume they want seconds
        return int(time_str)
    
    timespan_regex = re.compile(r"((?P<hours>\d+?)h)?((?P<minutes>\d+?)m)?((?P<seconds>\d+?)s)?")
    parts = timespan_regex.match(time_str)
    if not parts:
        raise ValueError("Invalid time span format. Valid formats: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc")
    parts = parts.groupdict() # {'hours': '3', 'minutes': '30', 'seconds': '10'}, {'hours': None, 'minutes': None, 'seconds': '20'}
    time_params = {name:int(value) for name, value in parts.items() if value}
    if not time_params:
        raise ValueError("Invalid time span format. Valid formats: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc")
    return int(timedelta(**time_params).total_seconds())


async def main():
    options = await parse_options()
    
    loglevel = options.loglevel.upper()
    if loglevel in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
        logfile = options.logfile
        if logfile:
            setup_logging(loglevel, logfile)
        else:
            setup_logging(loglevel, None)
    else:
        logger.warning("Invalid --loglevel. Valid values are: DEBUG/INFO/WARNING/ERROR/CRITICAL")
        sys.exit(1)
    
    aiotestfile = find_aiotestfile(options.aiotestfile)
    if not aiotestfile:
        logger.warning("Could not find any Aiotestfile! Ensure file ends in '.py' and see --help for available options.")
        sys.exit(1)
    
    num_users = options.num_users
    if not isinstance(num_users, int) or num_users <= 0:
        logger.warning(f"{options.num_users} mast be int type and >= 1")  
        sys.exit(1)
    
    rate = options.rate
    if not isinstance(rate, int) or rate <= 0 or rate > num_users:
        logger.warning(f"{options.rate} mast be int type and >= 1, and mast <= num_users")  
        sys.exit(1)
    
    user_classes, shape_class = load_aiotestfile(aiotestfile)
    if not user_classes:
        logger.warning("No User class found!")
        sys.exit(1)
    user_classes = list(user_classes.values())
    
    if os.name != "nt" and not options.master:
        try:
            import resource
            if resource.getrlimit(resource.RLIMIT_NOFILE)[0] < 10000:
                # Increasing the limit to 10000 within a running process should work on at least MacOS.
                # It does not work on all OS:es, but we should be no worse off for trying.
                resource.setrlimit(resource.RLIMIT_NOFILE, [10000, resource.RLIM_INFINITY])
        except BaseException:
            logger.warning(
                """
                System open file limit '10000' is below minimum setting '10000'. 
                It's not high enough for load testing, and the OS didn't allow aiotest to increase it by itself. 
                """
            )    
   
    if shape_class and (num_users or rate or options.run_time):
        logger.warning("aiotestfile contains the LoadUserShape subclass, which ignores arguments specified on the command line :users or rate or run_time.")
        options.run_time = None

    if options.show_users_wight:
        user_wight = {}
        for user in user_classes:
            user_wight[user.__name__] = user.weight
        pprint(user_wight)
        sys.exit(0)

    if options.run_time and not shape_class:
        if options.worker:
            logger.error("--run-time should be specified on the master node, and not on worker nodes")
            sys.exit(1)
        try:
            options.run_time = parse_timespan(options.run_time)
        except ValueError:
            logger.error("Valid --run-time formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
            sys.exit(1)
    
    options.buckets = sorted(set(options.buckets))
    options.buckets.append(float('inf'))
    stats_exporter.aiotest_response_times = Histogram(name='aiotest_response_times', documentation='aiotest response times', labelnames=["name", "method", "code"], buckets=tuple(options.buckets))
    
    if options.master:
        runners.global_runner = MasterRunner(user_classes, shape_class, options)
        while len(runners.global_runner.workers.ready) < options.expect_workers:
            logger.info(f"Waiting for workers to be ready, {len(runners.global_runner.workers.ready)} of {options.expect_workers} connected")
            await asyncio.sleep(2)     
    elif options.worker:
        runners.global_runner = WorkerRunner(user_classes, shape_class, options)   
    else:
        runners.global_runner = LocalRunner(user_classes, shape_class, options)
    
    await events.init.fire()
               
    try:
        if not options.worker:           
            async with asyncio.timeout(options.run_time):
                if shape_class:
                    runners.global_runner.start_shape()
                else:
                    await runners.global_runner.start(num_users, rate)
                            
                await asyncio.sleep(1)
                for task in runners.global_runner.tasks:
                    await task
        else:
            await asyncio.sleep(1)
            for task in runners.global_runner.tasks:
                await task 
    except (TimeoutError, KeyboardInterrupt, asyncio.CancelledError):
        await runners.global_runner.quit()
        await asyncio.sleep(0.2)
        if STATSERROR:
            for k, e in STATSERROR.items():
                logger.error(f"{k}:{e}")
            sys.exit(1)
        else:
            sys.exit(0)


def run():
    if sys.platform != "win32": 
        import uvloop
        uvloop.install()

    asyncio.run(main())







aiotest/main.py
hewei198711/Aiotest
# encoding: utf-8

import msgpack


class Message():
    def __init__(self, message_type, data, node_id):
        self.type = message_type
        self.data = data
        self.node_id = node_id
    
    def __repr__(self):
        return f"<Message {self.type}:{self.node_id}>"
    
    def serialize(self):
        "Pack object o and return packed bytes"
        return msgpack.packb((self.type, self.data, self.node_id))
    
    @classmethod
    def unserialize(cls, data):
        "Unpack an object"
        return cls(*msgpack.unpackb(data, raw=False, strict_map_key=False))

aiotest/rpc/protocol.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import asyncio
import traceback
from functools import wraps
import zmq
import zmq.asyncio as asynczmq
import zmq.error as zmqerr
import msgpack.exceptions as msgerr
from aiotest.exception import RPCError
from aiotest.rpc.protocol import Message
from aiotest.log import logger


def async_retry(func):
    "Retry 3 times"
    @wraps(func)
    async def wrapper(*args, **kwargs):
        delay = 0
        while delay <= 4:
            try:
                return await func(*args, **kwargs)
            except RPCError:
                delay += 2
                logger.error(f"Exception found on retry {delay/2}: --retry after {delay}")
                await asyncio.sleep(delay)
    return wrapper


class BaseSocket():
    def __init__(self, sock_type):
        context = asynczmq.Context()
        self.socket = context.socket(sock_type)
        self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1)
        self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30)

    @async_retry
    async def send(self, msg):
        try:
            await self.socket.send(msg.serialize(), zmq.NOBLOCK)
        except zmqerr.ZMQError:
            raise RPCError("ZMQ sent failure")

    @async_retry
    async def send_to_worker(self, msg):
        try:
            await self.socket.send_multipart([msg.node_id.encode(), msg.serialize()])
        except zmqerr.ZMQError:
            raise RPCError("ZMQ sent failure")
 
    async def recv(self):
        try:
            data = await self.socket.recv()
            msg = Message.unserialize(data)
        except msgerr.ExtraData:
            raise RPCError("ZMQ interrupted message")
        except zmqerr.ZMQError:
            raise RPCError("ZMQ network broken")
        return msg

    async def recv_from_worker(self):
        try:
            data = await self.socket.recv_multipart()
            addr = data[0].decode()
            msg = Message.unserialize(data[1])
        except (UnicodeDecodeError, msgerr.ExtraData):
            raise RPCError("ZMQ interrupted message")
        except zmqerr.ZMQError:
            raise RPCError("ZMQ interrupted message")
        return addr, msg
    
    def close(self):
        self.socket.close()


class Server(BaseSocket):
    def __init__(self, host, port):
        super().__init__(zmq.ROUTER)
        try:
            self.socket.bind(f"tcp://{host}:{port}")
            self.port = port
        except zmqerr.ZMQError as e:
            raise RPCError(f"Socket bind failure: {e}")


class Client(BaseSocket):
    def __init__(self, host, port, node_id):
        super().__init__(zmq.DEALER)
        self.socket.setsockopt(zmq.IDENTITY, node_id.encode())
        self.socket.connect(f"tcp://{host}:{port}")
aiotest/rpc/zmqrpc.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import socket
import asyncio
from uuid import uuid4
from aiotest.log import logger
from aiotest import events
from aiotest.rpc.protocol import Message
from aiotest.rpc.zmqrpc import Server, Client
from aiotest.exception import RPCError
from aiotest.stats_exporter import exporter_cpu_usage, exporter_user_count, prometheus_server

global_runner = None

HEARTBEAT_INTERVAL = 1 # worker 心跳监控时间间隔
HEARTBEAT_LIVENESS = 3 # 心跳
FALLBACK_INTERVAL = 5 # worker信息接收失败后，重新获取信息间隔
STATE_INIT, STATE_STARTING, STATE_RUNNING, STATE_STOPPED, STATE_MISSING = [
    "ready",
    "starting",
    "running", 
    "stopped", 
    "missing"
]

    
class Runner:
    def __init__(self, user_classes, shape_class, options):
        self.user_classes = user_classes
        self.shape_class = shape_class
        self.options = options
        self.cpu_usage = 0
        self.state = STATE_INIT
        self.user_instances = []
        self.tasks = []
        self.is_quit = False
        
    def __del__(self):
        for task in self.tasks:
            if not task.done():
                task.cancel()
               
    @property
    def user_count(self):
        "Returns the number of running user tasks, a user equals a task"
        return len([instances for instances in self.user_instances if not instances.task.done()])
    
    def update_state(self, new_state):
        logger.debug(f"Updating state to {new_state}, old state was {self.state}")   
        self.state = new_state
                       
    def weight_users(self, amount):
        """
        Distributes the amount of users for each User-class according to it's weight returns a list "bucket" with the weighted users
        """
        weight_sum = 0
        for user in self.user_classes:
            if self.options.host:
                user.host = self.options.host
            try:
                if isinstance(user.weight, int) and user.weight >= 1:
                    weight_sum += user.weight
                else:
                    raise ValueError("weigth value must be an int type and >= 1")
            except KeyError:
                raise KeyError("Userclass must have weight attribute")
            
        residuals = {}
        bucket = []
        for user in self.user_classes:
            # create users depending on weight
            percent = user.weight / weight_sum
            num_users = round(amount * percent)
            residuals[user] = amount * percent - num_users
            
            bucket.extend([user for i in range(num_users)])
            
        if len(bucket) < amount:
            for user in [l for l, r in sorted(residuals.items(), key=lambda x: x[1], reverse=True)][:amount-len(bucket)]:
                bucket.append(user)
        elif len(bucket) > amount:
            for user in [l for l, r in sorted(residuals.items(), key=lambda x: x[1])][:len(bucket)-amount]:
                bucket.remove(user)
        # [<class 'User01'>, <class 'User02'>,...]    
        return bucket
    
    async def start(self, user_count, rate):
        "Create user tasks for a load test master entry"
        if not isinstance(user_count, int) or user_count <= 0:
            logger.warning(f"{user_count} mast be int type and >= 1")
            sys.exit(1)
        if rate <= 0 or rate > user_count:
            logger.warning(f"{rate} mast > 0 and < user_count {user_count}") 
            sys.exit(1) 
        if rate % 1 != 0:
            logger.warning(f"{rate} rate fractional part is't 0") 
        
        # Dynamically changing the user count
        if self.state in [STATE_STARTING, STATE_RUNNING]:
            logger.debug(f"Updating running test with {user_count} users, {rate:.2f} rate.")
            self.update_state(STATE_STARTING)
            if self.user_count > user_count:
                # stop some users
                stop_count = self.user_count - user_count
                await self.stop_users(stop_count, rate)
                self.update_state(STATE_RUNNING)
            elif self.user_count < user_count:
                # start some users
                start_count = user_count - self.user_count
                await self.start_users(start_count, rate)
                self.update_state(STATE_RUNNING)
            else:
                self.update_state(STATE_RUNNING)
        elif self.state == STATE_INIT:
            await self.start_users(user_count, rate)
            self.update_state(STATE_RUNNING)
            await events.start_complete.fire(user_count=self.user_count, runner=self)
        else:
            logger.error(f"runner state is {self.state}")
            sys.exit(1)
    
    async def start_users(self, user_count, rate):  
        "Create a specified number of user tasks, a user equals a task"         
        bucket = self.weight_users(user_count) # [<class 'User01'>, <class 'User02'>,...]  
        if self.state == STATE_INIT:
            self.update_state(STATE_STARTING)
            
        existing_count = self.user_count
        logger.info(f"starting {len(bucket)} users at the rate {rate} users/s, ({existing_count} users already running)...")
        start_count = dict((u.__name__, 0) for u in self.user_classes) # {'User01': 0, 'User02': 0...}
        
        sleep_time = 1 / rate
        for i in range(len(bucket)):
            user_class = bucket.pop()
            start_count[user_class.__name__] += 1
            
            new_user = user_class()
            new_user.start_user()
            self.user_instances.append(new_user)
            
            await asyncio.sleep(sleep_time) 
            if self.user_count % 10 == 0:
                logger.debug(f"{self.user_count} users started")
            
        if not bucket:
            logger.info(f"All users started: {', '.join([f'{name}:{count}' for name, count in start_count.items()])}")
        else:
            logger.error(f"{bucket}  have user_class don't started")
            sys.exit(1)
        
    async def stop_users(self, user_count, rate):  
        "Cancels a specified number of user tasks"
        bucket = self.weight_users(user_count) # [<class 'User01'>, <class 'User02'>,...]  
        user_count = len(bucket)
        stop_users = []
        for instances in self.user_instances:
            if not instances.task.done():
                for user_class in bucket:
                    if isinstance(instances, user_class):
                        stop_users.append(instances)
                        bucket.remove(user_class)
                        break
                    
        if rate >= user_count:
            sleep_time = 0
            logger.info(f"Stopping {user_count} users immediately")
        else:
            sleep_time = 1 / rate
            logger.info(f"Stoping {user_count} users at rate of {rate} users/s")

        for instances in stop_users:
            await instances.stop_user()
            await asyncio.sleep(sleep_time)

        self.user_instances = [instances for instances in self.user_instances if not instances.task.done()]
        stop_users = [instances for instances in stop_users if not instances.task.done()]
        if stop_users:
            logger.warning(f"There are still user tasks uncancelled: {len(stop_users)}")
        
    async def stop(self):
        "Cancel all user tasks"
        logger.debug("Stopping all users")
        await self.stop_users(self.user_count, self.user_count)
        await asyncio.sleep(0.5)
        logger.debug(f"all user task is done: {all([instances.task.done() for instances in self.user_instances])}")
        for task in self.tasks:
            if task.get_name() == "users_tasks" and not task.done():
                task.cancel()
                self.tasks.remove(task)
                break
        self.update_state(STATE_STOPPED)

    async def quit(self):
        "Exit the load test and cancel all runner tasks"
        if not self.is_quit:
            self.is_quit = True
            await self.stop()
            await asyncio.sleep(0.5)
            await events.quitting.fire()
            for task in self.tasks:
                if not task.done() and task != asyncio.current_task():
                    task.cancel()
            await asyncio.sleep(0.1)
            self.tasks = [task for task in self.tasks if not task.done()]
            logger.debug(f"runner's tasks is done: {len(self.tasks) == 1}")
      
    def start_shape(self):
        "Create a load test policy task"
        shape_task = asyncio.create_task(self.shape_run(), name="shape_task")
        self.tasks.append(shape_task)
    
    async def shape_run(self):
        "Execute the specified load test policy"
        logger.info("Shape starting")
        shape_last = None
        while True:
            shape_new = self.shape_class.tick()
            if shape_new is None:
                logger.info("Shape test stopping")
                await self.quit()
                return
            elif shape_last == shape_new:
                await asyncio.sleep(1)
            else:
                logger.debug(shape_new)
                user_count, rate = shape_new
                logger.info(f"Shape test updating to {user_count} users at {rate} rate")
                await self.start(user_count=user_count, rate=rate)
                shape_last = shape_new
 

class LocalRunner(Runner):
    def __init__(self, user_classes, shape_class, options):
        super().__init__(user_classes, shape_class, options)
        user_count_task = asyncio.create_task(exporter_user_count(self), name="user_count_task")
        prometheus_task = asyncio.create_task(prometheus_server(self.options.prometheus_port), name="prometheus")
        monitor_cpu_task = asyncio.create_task(exporter_cpu_usage(self), name="monitor_cpu")
        
        self.tasks.append(user_count_task)
        self.tasks.append(prometheus_task)
        self.tasks.append(monitor_cpu_task)
    
    async def start(self, user_count, rate):
        if rate > 100:
            logger.warning("Your selected rate is very high (>100), and this is known to sometimes cause issues. Do you really need to ramp up that fast?")
        if self.state == STATE_INIT:
            await events.test_start.fire(runner=self)
        for task in asyncio.all_tasks():
            if task.get_name() == "users_tasks" and not task.done():
                # cancel existing task(name='users_tasks') before we start a new one
                task.cancel()
                break
        users_tasks = asyncio.create_task(super().start(user_count, rate), name="users_tasks")
        self.tasks.append(users_tasks)
    
    async def stop(self):
        # if self.state == STATE_STOPPED:
        #     return
        await super().stop()
        await events.test_stop.fire(runner=self)


class DistributedRunner(Runner):
    def __init__(self, user_classes, shape_class, options):
        super().__init__(user_classes, shape_class, options)        
        self.master_host = options.master_host
        self.master_port = options.master_port
        self.master_bind_host = options.master_bind_host
        self.master_bind_port = options.master_bind_port
        self.heartbeat_liveness = HEARTBEAT_LIVENESS
        self.heartbeat_interval = HEARTBEAT_INTERVAL
        self.connection_broken = False # 连接断了
  
  
class WorkerNode:
    def __init__(self, id, state=STATE_INIT, heartbeat_liveness=HEARTBEAT_LIVENESS):
        self.id = id
        self.state = state
        self.user_count = 0
        self.heartbeat = heartbeat_liveness
        self.cpu_usage = 0


class MasterRunner(DistributedRunner):
    """
    Runner used to run distributed load tests across multiple processes and/or machines.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        class WorkerNodesDict(dict):
            def get_by_state(self, state):
                return [c for c in self.values() if c.state == state]
            
            @property
            def all(self):
                return self.values()
            
            @property
            def ready(self):
                return self.get_by_state(STATE_INIT)
            @property
            def starting(self):
                return self.get_by_state(STATE_STARTING)
            
            @property
            def running(self):
                return self.get_by_state(STATE_RUNNING)
            
            @property
            def missing(self):
                return self.get_by_state(STATE_MISSING)
        
        self.workers = WorkerNodesDict()
        self.server = Server(self.master_bind_host, self.master_bind_port)
        
        worekr_heartbeat_task = asyncio.create_task(self.worker_heartbeat(), name="worekr_heartbeat")
        worekr_listener_task = asyncio.create_task(self.worekr_listener(), name="worekr_listener")
        prometheus_task = asyncio.create_task(prometheus_server(self.options.prometheus_port), name="prometheus")
        monitor_cpu_task = asyncio.create_task(exporter_cpu_usage(self), name="monitor_cpu")
        
        self.tasks.append(worekr_heartbeat_task)
        self.tasks.append(worekr_listener_task)
        self.tasks.append(prometheus_task)
        self.tasks.append(monitor_cpu_task)
            
    @property
    def user_count(self):
        return sum([c.user_count for c in self.workers.values()])

    @property
    def worker_count(self):
        return len(self.workers.ready) + len(self.workers.starting) + len(self.workers.running)
   
    async def start(self, user_count, rate):
        num_workers = self.worker_count
        if not num_workers:
            logger.warning("You are running in distributed mode but have no worker servers connected. Please connect workers prior to swarming.")
            return
        
        worker_host = self.options.host
        worker_user_count = user_count // (num_workers or 1) # 每个worker分配user_count
        worker_rate = rate / (num_workers or 1) # 每个worker生成user速率
        remainig = user_count % num_workers # 剩下的user_count
        
        logger.info(f"Sending jobs of {worker_user_count} users and {worker_rate:.2f} rate to {num_workers} ready workers")
        if worker_rate > 100:
            logger.warning("Your selected rate is very high (>100/worker), and this is known to sometimes cause issues. Do you really need to ramp up that fast?")
        if self.state == STATE_INIT:
            await events.test_start.fire(runner=self)
            self.update_state(STATE_STARTING)
        for worker in (self.workers.ready +self.workers.starting +self.workers.running):
            data = {
                "user_count": worker_user_count,
                "rate": worker_rate,
                "host": worker_host
            }
            if remainig > 0:
                data["user_count"] += 1
                remainig -= 1
                
            logger.debug(f"Sending start users message to worker {worker.id}")
            await self.server.send_to_worker(Message("start", data, worker.id))

    async def stop(self):
        logger.debug("Stopping...")  
        # 避免直到tick（）发出另一个变化信号时才停止
        if self.shape_class:
            self.tactics_current = None
        for worker in self.workers.all:
            logger.debug(f"Sending stop message to worker {worker.id}")
            await self.server.send_to_worker(Message("stop", None, worker.id))

        await events.test_stop.fire(runner=self)
        self.update_state(STATE_STOPPED)

    async def quit(self):
        logger.debug("Quitting...")
        if not self.is_quit:
            self.is_quit = True
            await self.stop()
            await asyncio.sleep(0.5)
                
            for worker in self.workers.all:
                logger.debug(f"Sending quit message to worker {worker.id}")
                await self.server.send_to_worker(Message("quit", None, worker.id))
            
            await asyncio.sleep(0.5) # wait for final stats report from all workers
            await events.quitting.fire()
            for task in self.tasks:
                if not task.done():
                    task.cancel()
            self.server.close()
    
    async def worker_heartbeat(self):
        while True:
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            if self.connection_broken:
                await self.reset_connection()
                continue
            for worker in self.workers.all:
                if worker.heartbeat < 0 and worker.state != STATE_MISSING:
                    logger.warning(f"Worker {worker.id} failed to send heartbeat, setting state to missing.")
                    worker.state = STATE_MISSING
                    worker.user_count = 0
                    if self.worker_count - len(self.workers.missing) <= 0:
                        logger.warning("The last worker went missing, stopping test.")
                        await self.stop()
                else:
                    worker.heartbeat -= 1
                    
    async def reset_connection(self):
        logger.info("Reset connection to worker")
        try:
            self.server.close()
            self.server = Server(self.master_bind_host, self.master_bind_port)
        except RPCError as e:
            logger.error(f"Temporary failure when resetting connection: {e}, will retry later.")

    async def worekr_listener(self): 
        while True:
            try:
                worker_id, msg = await self.server.recv_from_worker()
                msg.node_id = worker_id            
            except RPCError as e:
                logger.error(f"RPCError found when receiving from worker: {e}")
                self.connection_broken = True
                await asyncio.sleep(FALLBACK_INTERVAL)
                continue
            self.connection_broken = False
            if msg.type == "heartbeat":
                if worker_id in self.workers:
                    c = self.workers[worker_id]
                    c.heartbeat = HEARTBEAT_LIVENESS
                    c.state = msg.data["state"]
                    c.cpu_usage = msg.data["cpu_usage"]
                    if c.cpu_usage >= 90:
                        logger.warning(f"Worker {worker_id} exceeded CPU threshold")
            elif msg.type == "stats":
                if worker_id not in self.workers:
                    logger.warning(f"Discarded report from unrecognized worker {worker_id}")
                else:
                    self.workers[worker_id].user_count = msg.data["user_count"]
                await events.worker_report.fire(runner=self, worker_id=worker_id, data=msg.data)
            elif msg.type == "error":
                await events.user_error.fire(runner=self, error=msg.data["error"])   
            elif msg.type == "ready":
                self.workers[worker_id] = WorkerNode(id=worker_id)
                logger.info(f"Worker {worker_id} reported as ready. Currently {self.worker_count} workers ready to swarm.")
                if self.state in [STATE_STARTING, STATE_RUNNING]:
                    # balance the load distribution when new worker joins
                    self.start(self.options.user_count, self.options.rate)
            elif msg.type == "starting":
                self.workers[worker_id].state = STATE_STARTING
            elif msg.type == "start_complete":
                self.workers[worker_id].state = STATE_RUNNING
                self.workers[worker_id].user_count = msg.data["user_count"]
                if len(self.workers.running) == len(self.workers.all):
                    self.update_state(STATE_RUNNING)
                    await events.start_complete.fire(user_count = self.user_count, runner=self)
            elif msg.type == "stopped":
                if worker_id in self.workers:
                    del self.workers[worker_id]
                    logger.info(f"Removing {worker_id} worker from running workers")
            elif msg.type == "quitted":
                if worker_id in self.workers:
                    del self.workers[worker_id]
                    logger.info(f"Worker {worker_id} quitted. Currently {self.worker_count} workers connected.")


class WorkerRunner(DistributedRunner):
    "Runner used to run distributed load tests across multiple processes and/or machines."
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.worker_state = STATE_INIT
        # 主机ip地址 + 随机码
        self.worker_id = socket.gethostbyname(socket.gethostname()) + "_" + uuid4().hex
        self.worker = Client(self.master_host, self.master_port, self.worker_id)
        
        heartbeat_task = asyncio.create_task(self.heartbeat(), name="heartbeat")
        worker_run_task = asyncio.create_task(self.worker_run(), name="worker_run")
        monitor_cpu_task = asyncio.create_task(exporter_cpu_usage(self), name="monitor_cpu")
        
        self.tasks.append(heartbeat_task)
        self.tasks.append(worker_run_task)
        self.tasks.append(monitor_cpu_task)
        
        events.start_complete += self.on_start_complete
        events.quitting += self.on_quitting
 
    async def on_start_complete(self, user_count, runner):
        await self.worker.send(Message("start_complete", {"user_count":user_count}, self.worker_id))
        self.worker_state = STATE_RUNNING

    async def on_quitting(self):
        await self.worker.send(Message("quitted", None, self.worker_id)) 
    
    async def heartbeat(self):
        while True:
            try:
                await self.worker.send(
                    Message(
                        "heartbeat",
                        {"state": self.worker_state, "cpu_usage": self.cpu_usage},
                        self.worker_id
                    )
                )
            except RPCError as e:
                logger.error(f"RPCError found when sending heartbeat: {e}")
                self.reset_connection()
            finally:
                await asyncio.sleep(HEARTBEAT_INTERVAL)
    
    async def reset_connection(self):
        logger.info("Reset connection to master")
        try:
            self.worker.close()
            self.worker = Client(self.master_host, self.master_port, self.worker_id)
        except RPCError as e:
            logger.error(f"Temporary failure when resetting connection: {e}, will retry later.")
        
    async def worker_run(self):
        await self.worker.send(Message("ready", None, self.worker_id))
        while True:
            try:
                msg = await self.worker.recv()             
            except RPCError as e:
                logger.error(f"RPCError found when receiving from master: {e}")
                continue
            if msg.type == "start":
                self.worker_state = STATE_STARTING
                await self.worker.send(Message("starting", None, self.worker_id))
                job = msg.data
                self.options.host = job["host"]
                for task in asyncio.all_tasks():
                    if task.get_name() == "users_tasks" and not task.done():
                        # cancel existing task(name='start_user') before we start a new one
                        task.cancel()
                        break
                users_tasks = asyncio.create_task(self.start(job["user_count"], job["rate"]), name="users_tasks")
                self.tasks.append(users_tasks)
            elif msg.type == "stop":
                await self.stop()
                await self.worker.send(Message("stopped", None, self.worker_id))
                self.worker_state = STATE_STOPPED
                await self.worker.send(Message("ready", None, self.worker_id))
                self.worker_state = STATE_INIT
            elif msg.type == "quit":
                logger.info("Got quit message from master, shutting down...")
                await self.quit()
                break
            
    async def quit(self):
        "Exit the load test and cancel all runner tasks"
        await super().quit()
        if self.worker:
            self.worker.close()

aiotest/runners.py
hewei198711/Aiotest
# encoding: utf-8

from timeit import default_timer
from abc import ABCMeta, abstractmethod


class LoadUserShape(metaclass=ABCMeta):
    """
    A simple load test shape class used to control the shape of load generated
    during a load test.
    """
    def __init__(self):
        self.start_time = default_timer()

    def reset_time(self):
        "Resets start time back to 0"
        self.start_time = default_timer()
    
    def get_run_time(self):
        "Calculates run time in seconds of the load user"
        return default_timer() - self.start_time
        
    @abstractmethod
    def tick(self):
        """
        Returns a tuple with 2 elements to control the running load user:
            user_count -- Total user count
            rate -- Number of users to start/stop per second when changing number of users
        if 'None' is returned then the running load user will be stopped.
        
        """
        return None

aiotest/shape.py
hewei198711/Aiotest
# encoding: utf-8

import re
import psutil
import asyncio
import hashlib
from prometheus_client import Gauge, Counter, start_http_server
from aiotest import events
from aiotest.rpc.protocol import Message
from aiotest.log import logger

CPU_MONITOR_INTERVAL = 5 # CPU 监控时间间隔
USER_MONITOR_INTERVAL = 5 # USER_COUNT 监控时间间隔
STATSERROR = {} # 保存error

LABELS = ["name", "method", "code"]

aiotest_workers_user_count = Gauge(name='aiotest_workers_user_count', documentation='aiotest workers user count', labelnames=["node"])
aiotest_workers_cpu_usage = Gauge(name='aiotest_workers_cpu_usage', documentation='aiotest workers cpu usage', labelnames=["node"])
aiotest_response_times = None
aiotest_response_content_length = Gauge(name='aiotest_response_content_length', documentation='aiotest response content length', labelnames=LABELS)
aiotest_response_failure = Counter(name='aiotest_response_failure', documentation='aiotest response failure', labelnames=["name","method", "error"])
aiotest_user_error = Counter(name='aiotest_user_error', documentation='aiotest user error', labelnames=["error"])

code_regexp = re.compile(r"^\{\"code\":(\d{3})")


def parse_error(error):
    string_error = repr(error)
    if string_error.find("当前登录已失效，请重新登录") > 0:
        string_error = re.sub(r", \'time\': \'20\S+\'", "",string_error)
        if string_error.find("time") > 0:
            string_error = re.sub(r',\"time\":\"20\S+\"', "",string_error)
        return string_error

    target = "object at 0x"
    target_index = string_error.find(target)
    if target_index < 0:
        return string_error
    start = target_index + len(target) - 2
    end = string_error.find(">", start)
    if end < 0:
        return string_error
    hex_address = string_error[start:end]
    return string_error.replace(hex_address, "0x....")


async def on_request(runner, request_name, request_method, response_time, response_length, error):
    if not error:
        if type(runner).__name__ == "WorkerRunner":
            try:
                await runner.worker.send(
                    Message(
                        "stats", 
                        {
                            "request_name": request_name, 
                            "request_method": request_method, 
                            "response_time": response_time,
                            "response_length": response_length,
                            "error": error,
                            "user_count": runner.user_count
                        }, 
                        runner.worker_id
                    )
                )
            except:
                logger.error("Connection lost to master server. Aborting...")
        else:
            aiotest_response_times.labels(request_name, request_method, 200).observe(response_time)
            aiotest_response_content_length.labels(request_name, request_method, 200).set(response_length)  
    else:
        error = parse_error(error) 
        if type(runner).__name__ == "WorkerRunner":
            try:
                await runner.worker.send(
                    Message(
                        "stats", 
                        {
                            "request_name": request_name, 
                            "request_method": request_method, 
                            "response_time": response_time,
                            "response_length": response_length,
                            "error": error,
                            "user_count": runner.user_count
                        }, 
                        runner.worker_id
                    )
                )
            except:
                logger.error("Connection lost to master server. Aborting...")
        else:
            aiotest_response_times.labels(request_name, request_method, 400).observe(response_time)
            aiotest_response_content_length.labels(request_name, request_method, 400).set(response_length)
            aiotest_response_failure.labels(request_name, request_method, error).inc()
            key = create_key(error, request_name, request_method)
            if not STATSERROR.get(key):
                STATSERROR[key] = 1
                logger.error(f"{request_name}: {error}")
            else:
                STATSERROR[key] += 1

      
async def on_error(runner, error):
    error = parse_error(error)
    if type(runner).__name__ == "WorkerRunner":
        try:
            await runner.worker.send(Message("error", {"error": error}, runner.worker_id))
        except:
            logger.error("Connection lost to master server. Aborting...")
    else:
        aiotest_user_error.labels(error).inc()
        key = create_key(error)
        if not STATSERROR.get(key):
            STATSERROR[key] = 1
            logger.error(error)
        else:
            STATSERROR[key] += 1
        

async def on_worker_report(runner, worker_id, data):
    if worker_id not in runner.workers:
        logger.warning(f"Discarded report from unrecognized worker {worker_id}")
        return
    runner.workers[worker_id].user_count = data["user_count"]
    if not data["error"]:  
        aiotest_response_times.labels(data["request_name"], data["request_method"], 200).observe(data["response_time"])
        aiotest_response_content_length.labels(data["request_name"], data["request_method"], 200).set(data["response_length"])
    else:
        aiotest_response_times.labels(data["request_name"], data["request_method"], 400).observe(data["response_time"])
        aiotest_response_content_length.labels(data["request_name"], data["request_method"], 400).set(data["response_length"])
        aiotest_response_failure.labels(data["request_name"], data["request_method"], data["error"]).inc()
        key = create_key(data["error"], data["request_name"], data["request_method"])
        if not STATSERROR.get(key):
            STATSERROR[key] = 1
            logger.error(f'{data["request_name"]}: {data["error"]}')
        else:
            STATSERROR[key] += 1

    aiotest_workers_user_count.labels(worker_id).set(data["user_count"])         


async def exporter_cpu_usage(runner):
    process = psutil.Process()
    while True:
        cpu_usage = process.cpu_percent() # 当前进程CPU利用率的百分比
        if type(runner).__name__ == "MasterRunner":
            aiotest_workers_cpu_usage.labels("Master").set(cpu_usage)
            for id, worker in runner.workers.items():
                aiotest_workers_cpu_usage.labels(id).set(worker.cpu_usage)
            if cpu_usage >= 90:
                logger.warning(f"Master CPU usage {cpu_usage}!")
        elif type(runner).__name__ == "LocalRunner":
            aiotest_workers_cpu_usage.labels("Local").set(cpu_usage)  
            if cpu_usage >= 90:
                logger.warning(f"Local CPU usage {cpu_usage}!")
        else:
            runner.cpu_usage = cpu_usage
        await asyncio.sleep(CPU_MONITOR_INTERVAL)


async def exporter_user_count(runner):
    while True:
        aiotest_workers_user_count.labels("local").set(runner.user_count)
        await asyncio.sleep(USER_MONITOR_INTERVAL)

       
async def prometheus_server(port):
    start_http_server(port)


events.stats_request += on_request
events.user_error += on_error
events.worker_report += on_worker_report


def create_key(error, name=None, method=None):
    key = f"{method}.{name}.{error}"
    return hashlib.md5(key.encode('utf-8')).hexdigest()   


aiotest/stats_exporter.py
hewei198711/Aiotest
# encoding: utf-8

import asyncio
import base64
import random
import time
from io import BytesIO

import pytest
import allure
from asgiref.wsgi import WsgiToAsgi
from hypercorn.config import Config
from hypercorn.asyncio import serve
from flask import Flask, request, redirect, make_response, send_file, Response, stream_with_context
from aiotest import events
from aiotest.stats_exporter import on_request, on_error, on_worker_report
from aiotest.test.log import logger

app = Flask(__name__)
asgi_app = WsgiToAsgi(app)


@app.route("/")
async def get():
    await asyncio.sleep(0.1)
    return "asyncio sleep"

@app.route("/fast")
async def ultra_fast():
    return "This is an ultra fast response"

@app.route("/slow")
async def slow():
    delay = request.args.get("delay")
    if delay:
        await asyncio.sleep(float(delay))
    else:
        await asyncio.sleep(random.choice([0.5, 1, 1.5]))
    return "This is a slow response"

@app.route("/key")
async def keys():
    dic = {"key01": 1, "key02": 2}
    return dic

@app.route("/consistent")
async def consistent():
    await asyncio.sleep(0.2)
    return "This is a consistent response"

@app.route("/request_method", methods=["POST", "GET", "HEAD", "PUT", "DELETE", "PATCH"])
async def request_method():
    return request.method

@app.route("/request_header_test")
async def request_header_test():
    return request.headers["X-Header-Test"]

@app.route("/post", methods=["POST"])
@app.route("/put", methods=["PUT"])
async def manipulate():
    return str(request.form.get("arg", ""))

@app.route("/fail")
async def failed_request():
    return "This response failed", 500

@app.route("/status/204")
async def status_204():
    return "", 204

@app.route("/redirect", methods=["GET", "POST"])
async def do_redirect():
    delay = request.args.get("delay")
    if delay:
        await asyncio.sleep(float(delay))
    url = request.args.get("url", "/fast")
    return redirect(url)

@app.route("/basic_auth")
async def basic_auth():
    auth = base64.b64decode(request.headers.get("Authorization", "").replace("Basic ", "")).decode('utf-8')
    if auth == "aiotest:menace":
        return "Authorized"
    resp = make_response("401 Authorization Required", 401)
    resp.headers["WWW-Authenticate"] = 'Basic realm="aiotest"'
    return resp

@app.route("/no_content_length")
async def no_content_length():
    r = send_file(BytesIO("This response does not have content-length in the header".encode('utf-8')),
                  etag=False,
                  mimetype='text/plain')
    r.headers.remove("Content-Length")
    return r

@app.errorhandler(404)
async def not_found(error):
    return "Not Found", 404

@app.route("/streaming/<int:iterations>")
def streaming_response(iterations):
    def generate():
        yield "<html><body><h1>streaming response</h1>"
        for i in range(iterations):
            yield f"<span>{i}</span>\n"
            time.sleep(0.01)
        yield "</body></html>"
    return Response(stream_with_context(generate()), mimetype="text/html")

@app.route("/set_cookie", methods=["POST"])
async def set_cookie():
    response = make_response("ok")
    response.set_cookie(request.args.get("name"), request.args.get("value"))
    return response

@app.route("/get_cookie")
async def get_cookie():
    return make_response(request.cookies.get(request.args.get("name"), ""))

@pytest.fixture(scope="session")
def event_loop():
    loop = asyncio.new_event_loop()
    yield loop
    loop.close()

@allure.title("启动/关闭flask服务")
@pytest.fixture(scope="session", autouse=True)
async def web_server_up():
    config = Config()
    config.bind = ["localhost:8080"]
    flask_task = asyncio.create_task(serve(asgi_app, config), name="flask_task")
    events.stats_request -= on_request
    events.user_error -= on_error
    events.worker_report -= on_worker_report
    yield
    flask_task.cancel()
    await asyncio.sleep(0.2)

        





aiotest/test/conftest.py
hewei198711/Aiotest
# coding: utf-8

import sys

from loguru import logger

from aiotest.test.setting import LEVEL

logger.remove()
logger.add(sys.stderr, level=LEVEL)
logger.add("aiotest/test/log.log", level=LEVEL, retention="1 days")

aiotest/test/log.py
hewei198711/Aiotest
# coding:utf-8

BASE_URL = "http://127.0.0.1:8080"
P1 = "blocker"
P2 = "critical"
P3 = "normal"
LEVEL = "DEBUG"
aiotest/test/setting.py
hewei198711/Aiotest
# encoding: utf-8

import json
import pytest
import asyncio
import allure
from timeit import default_timer
from aiotest.client import ClientSession
from aiotest.stats_exporter import parse_error
from aiotest import events
from aiotest.test.setting import BASE_URL, P1, P2, P3
from aiotest.test.log import logger


request_dict = {}


async def on_request(runner, request_name, request_method, response_time, response_length, error):
    global request_dict
    request_dict = {}
    if error:
        error = parse_error(error)
    request_dict = {
        "request_name": request_name, 
        "request_method": request_method, 
        "response_time": response_time,
        "response_length": response_length,
        "error": error,
    }


@allure.title("add/remove events.stats_request")
@pytest.fixture(scope="module", autouse=True)
def events_handlers():
    events.stats_request += on_request
    yield
    events.stats_request -= on_request

@allure.title("init request_dict/error_dict")
@pytest.fixture(scope="function", autouse=True)
def init_request_and_error_dict():
    global request_dict
    request_dict = {}


@allure.feature("aiotest/client")
@allure.story("aiotest/client/ClientSession")
class TestClientSession:
        
    @classmethod
    @allure.title("start ClientSession")
    def setup_class(cls):
        cls.session = ClientSession(BASE_URL)

    @classmethod
    @allure.title("close ClientSession")
    def terdown_class(cls):
        cls.session.close()
                
    @allure.severity(P1)
    @allure.title("get method")
    async def test_01(self):
        start_time = default_timer()
        async with self.session.get("/", timeout=8) as resp: 
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
            assert text == "asyncio sleep"
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"] == None  
    
    @allure.severity(P1)
    @allure.title("streaming")
    async def test_02(self): 
        start_time = default_timer()
        async with self.session.get("/streaming/30") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/streaming/30"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)  
        assert request_dict["error"] == None   

    @allure.severity(P1)
    @allure.title("redirect")
    async def test_03(self):
        
        url = "/redirect?url=/redirect?delay=0.2"
        start_time = default_timer()
        async with self.session.get(url) as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/fast"
        assert request_dict["request_method"] == "GET"
        assert 205 < abs(request_dict["response_time"] - times) > 200
        assert request_dict["response_length"] == len(text)  
        assert request_dict["error"] == None   
                        
    @allure.severity(P2)
    @allure.title("redirect and name")
    async def test_04(self):
        start_time = default_timer()        
        async with self.session.post("/redirect", name="/redirect") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/redirect"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 10
        assert request_dict["response_length"] == len(text)  
        assert request_dict["error"] == None  
            
    @allure.severity(P1)
    @allure.title("cookie")
    async def test_05(self):
        start_time = default_timer()   
        async with self.session.post("/set_cookie?name=testcookie&value=1337") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/set_cookie"
        assert request_dict["request_method"] == "POST"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)  
        assert request_dict["error"] == None
        
        start_time = default_timer()   
        async with self.session.get("/get_cookie?name=testcookie") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/get_cookie"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"] == None
    
    @allure.severity(P2)
    @allure.title("head method")
    async def test_06(self):
        start_time = default_timer()   
        async with self.session.head("/request_method") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/request_method"
        assert request_dict["request_method"] == "HEAD"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == 4 
        assert request_dict["error"] == None
        
    @allure.severity(P2)
    @allure.title("delete method")
    async def test_07(self):
        start_time = default_timer()   
        async with self.session.delete("/request_method") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/request_method"
        assert request_dict["request_method"] == "DELETE"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)  
        assert request_dict["error"] == None

    @allure.severity(P2)
    @allure.title("all method")
    async def test_08(self):
        start_time = default_timer()   
        async with self.session.options("/request_method") as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/request_method"
        assert request_dict["request_method"] == "OPTIONS"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"] == None
        assert set(["OPTIONS", "DELETE", "PUT", "GET", "POST", "HEAD", "PATCH"]) == set(resp.headers["allow"].split(", "))
            
    @allure.severity(P1)
    @allure.title("ClientResponseError")
    async def test_09(self):
        start_time = default_timer()   
        async with self.session.get('/wrong_url/01') as resp:
            assert resp.status == 404
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/wrong_url/01"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"].find("ClientResponseError") >= 0

    @allure.severity(P1)
    @allure.title("TimeoutError")
    async def test_10(self):
        with pytest.raises(TimeoutError):
            async with self.session.get("/consistent", timeout=0.1) as resp:
                text = await resp.text()
                 
    @allure.severity(P2)
    @allure.title("ClientResponseError")
    async def test_11(self):
        start_time = default_timer()   
        async with self.session.get('/wrong_url/01') as resp:
            assert resp.status == 404
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/wrong_url/01"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"].find("ClientResponseError") >= 0

    @allure.severity(P2)
    @allure.title("500")
    async def test_11(self):
        start_time = default_timer()   
        async with self.session.get('/fail') as resp:
            assert resp.status == 500
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
        assert request_dict["request_name"] == "/fail"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"].find("ClientResponseError") >= 0

    @allure.severity(P2)
    @allure.title("TypeError")
    async def test_12(self):
        start_time = default_timer()   
        async with self.session.get('/') as resp:
            assert resp.status == 200
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
            text["id"]
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"].find("TypeError") >= 0

    @allure.severity(P2)
    @allure.title("KeyError")
    async def test_13(self):
        start_time = default_timer()   
        async with self.session.get('/key') as resp:
            assert resp.status == 200
            data = await resp.json()
            times = int((default_timer() - start_time) * 1000)
            data["key03"]
        assert request_dict["request_name"] == "/key"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["error"].find("KeyError") >= 0

    @allure.severity(P2)
    @allure.title("AssertionError")
    async def test_14(self):
        start_time = default_timer()   
        async with self.session.get('/key') as resp:
            assert resp.status == 200
            data = await resp.json()
            times = int((default_timer() - start_time) * 1000)
            assert data["key01"] == 3
        assert request_dict["request_name"] == "/key"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["error"].find("AssertionError") >= 0

    @allure.severity(P2)
    @allure.title("success")
    async def test_15(self):
        start_time = default_timer()   
        async with self.session.get('/fast') as resp:
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
            if resp.status == 200:
                resp.success()
        assert request_dict["request_name"] == "/fast"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"] is None

    @allure.severity(P2)
    @allure.title("failure")
    async def test_16(self):
        start_time = default_timer()   
        async with self.session.get('/fast') as resp:
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
            if resp.status == 200:
                resp.failure("my is a failure")
        assert request_dict["request_name"] == "/fast"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"].find("CatchResponseError('my is a failure')") >= 0

    @allure.severity(P2)
    @allure.title("success error")
    async def test_17(self):
        start_time = default_timer()   
        async with self.session.get('/fail') as resp:
            text = await resp.text()
            times = int((default_timer() - start_time) * 1000)
            resp.success()
        assert request_dict["request_name"] == "/fail"
        assert request_dict["request_method"] == "GET"
        assert abs(request_dict["response_time"] - times) < 5
        assert request_dict["response_length"] == len(text)
        assert request_dict["error"] is None



aiotest/test/test_client.py
hewei198711/Aiotest
# encoding: utf-8

import os
import asyncio
import pytest
import allure
import aiohttp
from aiotest.test.log import logger
from aiotest.test.setting import P1, P2, P3
from aiotest.rpc.protocol import Message
from aiotest import User, AsyncHttpUser, LoadUserShape
from aiotest.main import parse_options, find_aiotestfile, load_aiotestfile, is_shape_class, is_user_class


@allure.feature("rpc")
@allure.story("main/parse_options")
class TestArgparse:
    
    @allure.severity(P1)
    @allure.title("parse options")
    async def test_01(self):
        args = await parse_options(args=[
            "-f", "aiotestfile",
            "-u", "100",
            "-r", "10",
        ])
        assert args.host == ""
        assert args.master == False
        assert args.aiotestfile == "aiotestfile"
        assert args.num_users == 100
        assert args.rate == 10
        assert args.loglevel == "INFO"
        assert args.logfile == None
        assert args.show_users_wight == False

        
@allure.feature("main")
@allure.story("main/aiotestfile")
class TestAiotestfile:
    
    mock_user_file_content = """
from aiotest.user import AsyncHttpUser
from aiotest.shape import LoadUserShape


class TestUser(AsyncHttpUser):
    host = "http://127.0.0.1:8080"
    
    async def test_func(self):
        pass

class TestUser02(AsyncHttpUser):
    host = "http://127.0.0.1:8080"
    
    async def test_func(self):
        pass

class UserSubclass03(AsyncHttpUser):
    host = "http://127.0.0.1:8080"
    
    async def func(self):
        pass
        
class UserSubclass():
    host = "http://localhost:8080"

class TestShape(LoadUserShape):
    def tick(self):
        pass
    """
    directory = os.path.dirname(os.path.abspath(__file__))
    filename = "aiotestfile" 

    @classmethod
    def setup_class(cls):
        cls.file_path = os.path.join(cls.directory, "aiotestfile.py")
        with open(cls.file_path, "w") as file:
            file.write(cls.mock_user_file_content)    

    @classmethod
    def teardown_class(cls):
        os.remove(cls.file_path)
            
    @allure.severity(P2)
    @allure.title("find aiotestfile")
    async def test_01(self):
        aiotestfile = find_aiotestfile("aiotest/test/aiotestfile.py")
        assert r"aiotest\test\aiotestfile.py" in aiotestfile

    @allure.severity(P3)
    @allure.title("find aiotestfile not .py")
    async def test_02(self):
        aiotestfile = find_aiotestfile("aiotest/test/aiotestfile")
        assert r"aiotest\test\aiotestfile.py" in aiotestfile

    @allure.severity(P3)
    @allure.title("find aiotestfile not .py")
    async def test_03(self):
        user_classes, shape_class = load_aiotestfile(self.file_path)
        assert len(user_classes) == 2
        logger.debug(user_classes)

    @allure.severity(P3)
    @allure.title("is user class")
    async def test_04(self):
        assert is_user_class(User) is False
        assert is_user_class(AsyncHttpUser) is False
        
        class TestUser01(AsyncHttpUser):
            pass
        
        class User02(AsyncHttpUser):
            pass
        
        assert is_user_class(TestUser01)
        assert is_user_class(User02) is False
        

    @allure.severity(P3)
    @allure.title("is shape class")
    async def test_05(self):
        assert is_shape_class(LoadUserShape) is False
        
        class TestShape(LoadUserShape):
            def tick(self):
                pass

        assert is_shape_class(TestShape)

    @allure.severity(P3)
    @allure.title("load aiotestfile from relative path")
    async def test_06(self):
        user_classes, shape_class = load_aiotestfile(os.path.join("./aiotest/test/", "aiotestfile.py"))
        assert "TestUser" in user_classes
        assert "TestUser02" in user_classes

    @allure.severity(P3)
    @allure.title("load aiotestfile from absolute path")
    async def test_06(self):
        user_classes, shape_class = load_aiotestfile(self.file_path)
        assert "TestUser" in user_classes
        assert "TestUser02" in user_classes


aiotest/test/test_main.py
hewei198711/Aiotest
# encoding: utf-8

import pytest
import allure
from aiotest.test.setting import P1, P2, P3
from aiotest.rpc.protocol import Message
from aiotest.rpc import zmqrpc
from aiotest.test.log import logger

@allure.feature("rpc")
@allure.story("rpc/msgpack")
class TestMsgpack:
    
    @allure.severity(P1)
    @allure.title("message serialize/unserialize")
    async def test_01(self):
        
        msg = Message(
            "stats", 
            {"name": "/getname", "method": "GET", "response_time": 960, "response_length": 56, "user_count": 1000}, 
            "127.0.0.1_asdfzxcvfd"
        )
        
        packb = msg.serialize()
        rebuilt = Message.unserialize(packb)

        assert msg.type == rebuilt.type
        assert msg.data == rebuilt.data
        assert msg.node_id == rebuilt.node_id
                        

@allure.feature("rpc")
@allure.story("rpc/zmq")
class TestZmq:
    
    @classmethod
    @allure.title("Start ZMQ")
    def setup_class(cls):
        cls.server = zmqrpc.Server("127.0.0.1", 5557)
        cls.client = zmqrpc.Client("127.0.0.1", 5557, "worker_01")
    
    @classmethod
    @allure.title("Close ZMQ")    
    def teardown_class(cls):
        cls.server.socket.close()
        cls.client.socket.close()
    
    @allure.severity(P1)
    @allure.title("client send / server revc")
    async def test_01(self):
        await self.client.send(
            Message(
            "stats", 
            {"name": "/getname", "method": "GET", "response_time": 960, "response_length": 56, "user_count": 1000}, 
            "worker_01"
            )
        )
        addr, msg = await self.server.recv_from_worker()
        assert addr == "worker_01"
        assert msg.type == "stats" and msg.node_id == "worker_01"
        assert msg.data == {"name": "/getname", "method": "GET", "response_time": 960, "response_length": 56, "user_count": 1000}

    @allure.severity(P1)
    @allure.title("server send / client revc")
    async def test_02(self):
        await self.server.send_to_worker(
            Message(
            "test", 
            "message", 
            "worker_01"
        ))
        msg = await self.client.recv()
        assert msg.type == "test"
        assert msg.data == "message"
        assert msg.node_id == "worker_01"
                
aiotest/test/test_rpc.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import traceback
import pytest
import asyncio
import allure
from asyncio import Queue

from aiotest.rpc.zmqrpc import Server, Client
from aiotest.rpc.protocol import Message
from aiotest.user import AsyncHttpUser
from aiotest.stats_exporter import parse_error
from aiotest import events
from aiotest import runners
from aiotest.runners import LocalRunner, MasterRunner, WorkerRunner
from aiotest.shape import LoadUserShape
from aiotest.runners import STATE_INIT, STATE_STARTING, STATE_RUNNING, STATE_MISSING, STATE_STOPPED
from aiotest.test.setting import BASE_URL, P1, P2, P3
from aiotest.test.log import logger

request_dict = {}
error_dict = {}
user_count_dict = {}
test_start = False
test_stop = False


async def on_request(runner, request_name, request_method, response_time, response_length, error):
    global request_dict
    request_dict = {}
    if error:
        error = parse_error(error)
    request_dict = {
        "request_name": request_name, 
        "request_method": request_method, 
        "response_time": response_time,
        "response_length": response_length,
        "error": error,
    }

async def on_error(runner, error):
    global error_dict
    error_dict = {}
    if error:
        error = parse_error(error)
    error_dict = {
        "error": error,
    }

async def on_worker_report(runner, worker_id, data):
    global request_dict, user_count_dict
    runner = runners.global_runner
    if worker_id not in runner.workers:
        logger.warning(f"Discarded report from unrecognized worker {worker_id}")
        return
    runner.workers[worker_id].user_count = data["user_count"]
    if data["error"]:  
        data["error"] = parse_error(data["error"])
    request_dict = {
        "request_name": data["request_name"], 
        "request_method": data["request_method"], 
        "response_time": data["response_time"],
        "response_length": data["response_length"],
        "error": data["error"],
    }

    user_count_dict[worker_id] = data["user_count"]

async def on_test_start(runner):
    global test_start
    test_start = True

async def on_test_stop(runner):
    global test_stop
    test_stop = True
    
class mocked_options:
    def __init__(self):
        self.host= BASE_URL
        self.master= False
        self.worker=False
        self.master_host= "localhost"
        self.master_port = 5557
        self.master_bind_host= "*"
        self.master_bind_port= 5557
        self.heartbeat_liveness=3
        self.heartbeat_interval=1
        self.expect_workers= 1
        self.user_count= 2
        self.rate= 1
        self.show_users_wight= False
    
    def reset_stats(self):
        pass

class mocked_rpc:
    def __init__(self, *args, **kwargs):
        self.queue = Queue()
        self.outbox = []
    
    async def mocked_send(self, message):
        "master&worker send"
        await self.queue.put(message.serialize())
    
    async def recv(self):
        "worker recv"
        data = await self.queue.get()
        return Message.unserialize(data)
    
    async def send(self, message):
        "worker send"
        self.outbox.append(message)
    
    async def send_to_worker(self, message):
        "master send"
        self.outbox.append((message.node_id, message))
    
    async def recv_from_worker(self):
        "master recv"
        data = await self.queue.get()
        msg = Message.unserialize(data)
        return msg.node_id, msg


@allure.title("add/remove events.stats_request/user_error/worker_report")
@pytest.fixture(scope="module", autouse=True)
def events_handlers():
    events.stats_request += on_request
    events.user_error += on_error
    events.worker_report += on_worker_report
    events.test_start += on_test_start
    events.test_stop += on_test_stop
    yield
    events.stats_request -= on_request
    events.user_error -= on_error
    events.worker_report -= on_worker_report
    events.test_start -= on_test_start
    events.test_stop -= on_test_stop
 
@allure.title("mock rpc")
@pytest.fixture(scope="function", autouse=True)
def rpc(monkeypatch):
    rpc = mocked_rpc()
    monkeypatch.setattr(Server, "send_to_worker", rpc.send_to_worker)
    monkeypatch.setattr(Server, "recv_from_worker", rpc.recv_from_worker)
    monkeypatch.setattr(Client, "send", rpc.send)
    monkeypatch.setattr(Client, "recv", rpc.recv)
    yield rpc
    monkeypatch.delattr(Server, "send_to_worker", rpc.send_to_worker)
    monkeypatch.delattr(Server, "recv_from_worker", rpc.recv_from_worker)
    monkeypatch.delattr(Client, "send", rpc.send)
    monkeypatch.delattr(Client, "recv", rpc.recv)
  

@allure.title("init request_dict/error_dict/user_count_dict")
@pytest.fixture(scope="function", autouse=True)
async def init_request__error_user_dict(rpc):
    global request_dict, error_dict, user_count_dict, test_start, test_stop
    request_dict = {}
    error_dict = {}
    user_count_dict = {}
    test_start = False
    test_stop = False
    yield
    await runners.global_runner.quit()


@allure.feature("aiotest/runners")
@allure.story("aiotest/runners/LocalRunner")
class TestLocalRunner:
        
    @allure.severity(P2)
    @allure.title("init userclass's host")
    async def test_01(self):
        class MyUser(AsyncHttpUser):
            pass
        
        runner = LocalRunner([MyUser], None, mocked_options())
        runners.global_runner = runner
        runner.weight_users(1)
        assert MyUser.host == runner.options.host

    @allure.severity(P1)
    @allure.title("calculate userclass's weight(equal)")
    async def test_02(self):
        class MyUser(AsyncHttpUser):
            pass

        class MyUser02(AsyncHttpUser):
            pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        bucket = runner.weight_users(10)
        assert bucket.count(MyUser) == 5
        assert bucket.count(MyUser02) == 5

    @allure.severity(P2)
    @allure.title("calculate userclass's weight(not equal)")
    async def test_03(self):

        class MyUser(AsyncHttpUser):
            weight = 1

        class MyUser02(AsyncHttpUser):
            weight = 2
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        bucket = runner.weight_users(10)
        assert bucket.count(MyUser) == 3
        assert bucket.count(MyUser02) == 7
 
    @allure.severity(P2)
    @allure.title("calculate userclass's weight(extremum == 1)")
    async def test_04(self):

        class MyUser(AsyncHttpUser):
            weight = 1

        class MyUser02(AsyncHttpUser):
            weight = 2
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        bucket = runner.weight_users(1)
        assert len(bucket) == 1
        assert bucket.count(MyUser02) == 1
 
    @allure.severity(P1)
    @allure.title("start user_count==4, then start user_count==2(userclasses)")
    async def test_05(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.5)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await runner.start(user_count=2, rate=2)
        await asyncio.sleep(1.5)
        assert runner.user_count == 2
        
    @allure.severity(P2)
    @allure.title("start user_count==2, then start user_count==4(userclasses)")
    async def test_06(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=2, rate=2)
        await asyncio.sleep(1.3)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 2
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.3)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        
    @allure.severity(P1)
    @allure.title("start_users user_count==2, then start_users user_count==4(userclasses)")
    async def test_07(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start_users(user_count=2, rate=2)
        await asyncio.sleep(1.2)
        assert runner.state == STATE_STARTING
        assert runner.user_count == 2
        await runner.start_users(user_count=2, rate=2)
        await asyncio.sleep(1.2)
        assert runner.state == STATE_STARTING
        assert runner.user_count == 4
        
    @allure.severity(P1)
    @allure.title("stop all user tasks")
    async def test_08(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.5)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await runner.stop()
        await asyncio.sleep(0.5)
        assert runner.state == STATE_STOPPED
        assert runner.user_count == 0
        
    @allure.severity(P1)
    @allure.title("stop_user portion user task")
    async def test_09(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.5)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await runner.stop_users(user_count=2, rate=2)
        await asyncio.sleep(1)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 2
        
    @allure.severity(P2)
    @allure.title("stop_user portion user task(rate < user_count)")
    async def test_10(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.5)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await runner.stop_users(user_count=2, rate=1)
        await asyncio.sleep(1)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 2
        
    @allure.severity(P1)
    @allure.title("quit")
    async def test_11(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class MyUser02(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser, MyUser02], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=4, rate=4)
        await asyncio.sleep(1.5)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await runner.quit()
        await asyncio.sleep(1)
        assert runner.state == STATE_STOPPED
        assert runner.user_count == 0
        for task in runner.tasks:
            assert task.done() is True
        
    @allure.severity(P1)
    @allure.title("load user shape")
    async def test_12(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        class TestShape(LoadUserShape):
            def tick(self):
                run_time = self.get_run_time()
                if run_time < 2:
                    return (2, 2)
                elif run_time < 4:
                    return (4, 4)
                elif run_time < 6:
                    return (2, 2)
                else:
                    return None
                
        runner = LocalRunner([MyUser], TestShape(), mocked_options())
        runners.global_runner = runner
        runner.start_shape()
        await asyncio.sleep(1.3)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 2
        await asyncio.sleep(1.7)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 4
        await asyncio.sleep(2.3)
        assert runner.state == STATE_RUNNING
        assert runner.user_count == 2

    @allure.severity(P1)
    @allure.title("events.test_start and events.test_stop")
    async def test_13(self):

        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass
                
        runner = LocalRunner([MyUser], None, mocked_options())
        runners.global_runner = runner
        await runner.start(user_count=2, rate=2)
        await asyncio.sleep(1.3)
        assert test_start is True
        assert test_stop is False
        await runner.stop()
        await asyncio.sleep(1.3)
        assert test_start is True
        assert test_stop is True

            
@allure.feature("aiotest/runners")
@allure.story("aiotest/runners/MasterRunner")
class TestMasterRunner:
        
    @allure.severity(P2)
    @allure.title("worker send 'ready' to master")
    async def test_01(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await asyncio.sleep(0.1)
        assert master.worker_count == 1
        assert "worker01" in master.workers
        
        await rpc.mocked_send(Message("ready", None, "worker02"))
        await rpc.mocked_send(Message("ready", None, "worker03"))
        await rpc.mocked_send(Message("ready", None, "worker04"))
        await asyncio.sleep(0.1)
        assert master.worker_count == 4
        
        await rpc.mocked_send(Message("quitted", None, "worker01"))
        await asyncio.sleep(0.1)
        assert master.worker_count == 3

    @allure.severity(P2)
    @allure.title("worker send 'starting' to master")
    async def test_02(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await asyncio.sleep(0.1)
        assert master.worker_count == 1
        assert "worker01" in master.workers
        
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0.1)
        assert master.workers["worker01"].state == STATE_STARTING

    @allure.severity(P2)
    @allure.title("worker send 'start_complete' to master")
    async def test_03(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        await rpc.mocked_send(Message("ready", None, "worker02"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(0.1)
        assert master.worker_count == 2
        assert master.user_count == 2
        assert master.workers["worker01"].user_count == 2
        
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker02"))
        await asyncio.sleep(0.1)
        assert master.worker_count == 2
        assert master.user_count == 4
        assert master.workers["worker01"].user_count == 2
        assert master.workers["worker02"].user_count == 2

    @allure.severity(P2)
    @allure.title("worker send 'stats' to master(not error)")
    async def test_04(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        stats = {
            "request_name": "/get", 
            "request_method": "GET", 
            "response_time": 230,
            "response_length": 60,
            "error": None,
            "user_count": 2
        }
        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(0.01)
        await rpc.mocked_send(Message("stats", stats, "worker01"))
        await asyncio.sleep(0.01)
        assert master.worker_count == 1
        assert master.user_count == 2
        assert master.workers["worker01"].user_count == 2
        user_count = stats.pop("user_count")
        assert request_dict == stats
        assert user_count_dict["worker01"] == user_count
        
    @allure.severity(P2)
    @allure.title("worker send 'stopped' to master")
    async def test_05(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(0.01)
        await rpc.mocked_send(Message("stopped", None, "worker01"))
        await asyncio.sleep(0.01)
        assert master.worker_count == 0
        assert master.user_count == 0
        await rpc.mocked_send(Message("ready", None, "worker01"))
        await asyncio.sleep(0.01)
        assert master.worker_count == 1
        assert master.user_count == 0
        
    @allure.severity(P3)
    @allure.title("worker send 'quit' to master")
    async def test_06(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(0.01)
        await rpc.mocked_send(Message("stopped", None, "worker01"))
        await asyncio.sleep(0.01)
        assert master.worker_count == 0
        assert master.user_count == 0
        await rpc.mocked_send(Message("ready", None, "worker01"))
        await asyncio.sleep(0.01)
        assert master.worker_count == 1
        assert master.user_count == 0
        await rpc.mocked_send(Message("quitted", None, "worker01"))
        await asyncio.sleep(0.1)
        assert master.worker_count == 0
        assert master.user_count == 0        

    @allure.severity(P3)
    @allure.title("worker send 'error' to master")
    async def test_07(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(0.01)
        await rpc.mocked_send(Message("error", {"error": "this is error"}, "worker01"))
        await asyncio.sleep(0.01)
        assert "this is error" in error_dict["error"]
     
    @allure.severity(P2)
    @allure.title("worker user count")
    async def test_08(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        stats = {
            "request_name": "/get", 
            "request_method": "GET", 
            "response_time": 230,
            "response_length": 60,
            "error": None,
            "user_count": 2
        }
        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("stats", stats, "worker01"))
        await asyncio.sleep(0.01)
        assert master.user_count == 2
     
    @allure.severity(P3)
    @allure.title("worker missing")
    async def test_09(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        await rpc.mocked_send(Message("ready", None, "worker01"))
        master = MasterRunner([MyUser], None, mocked_options())
        runners.global_runner = master
        await rpc.mocked_send(Message("starting", None, "worker01"))
        await asyncio.sleep(0)
        await rpc.mocked_send(Message("start_complete", {"user_count": 2}, "worker01"))
        await asyncio.sleep(6)   
        assert master.user_count == 0
        assert master.worker_count == 0
        assert master.workers["worker01"].state == STATE_MISSING
     

@allure.feature("aiotest/runners")
@allure.story("aiotest/runners/WorkerRunner")
class TestWorkerRunner:
        
    @allure.severity(P2)
    @allure.title("worker send 'ready' to master")
    async def test_01(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        worker = WorkerRunner([MyUser], None, mocked_options())
        runners.global_runner = worker
        await asyncio.sleep(0.1)
        assert "ready" in [msg.type for msg in rpc.outbox]

    @allure.severity(P1)
    @allure.title("worker recv 'start'")
    async def test_02(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        worker = WorkerRunner([MyUser], None, mocked_options())
        runners.global_runner = worker
        await asyncio.sleep(0.1)
        # master send
        await rpc.mocked_send(
            Message(
                "start", 
                {
                    "user_count": 2,
                    "rate": 2,
                    "host": BASE_URL,
                },
                rpc.outbox[0].node_id
            )
        )
        await asyncio.sleep(1)
        assert "starting" in [msg.type for msg in rpc.outbox]
        assert worker.user_count == 2
        await asyncio.sleep(0.1)
        assert worker.worker_state == STATE_RUNNING
        assert "start_complete" in [msg.type for msg in rpc.outbox]

    @allure.severity(P2)
    @allure.title("worker recv 'stop'")
    async def test_03(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        worker = WorkerRunner([MyUser], None, mocked_options())
        runners.global_runner = worker
        await asyncio.sleep(0.1)
        # master send
        await rpc.mocked_send(
            Message(
                "start", 
                {
                    "user_count": 2,
                    "rate": 2,
                    "host": BASE_URL,
                },
                rpc.outbox[0].node_id
            )
        )
        await asyncio.sleep(1)
        assert worker.user_count == 2
        # master send
        await rpc.mocked_send(Message("stop", None, rpc.outbox[0].node_id))
        await asyncio.sleep(1)
        assert "stopped" in [msg.type for msg in rpc.outbox]
        assert worker.user_count == 0
        assert worker.worker_state == STATE_INIT

    @allure.severity(P2)
    @allure.title("worker recv 'quit'")
    async def test_04(self, rpc):
        class MyUser(AsyncHttpUser):
            wait_time = 0
            async def test(self):
                pass

        worker = WorkerRunner([MyUser], None, mocked_options())
        runners.global_runner = worker
        await asyncio.sleep(0.1)
        # master send
        await rpc.mocked_send(
            Message(
                "start", 
                {
                    "user_count": 2,
                    "rate": 2,
                    "host": BASE_URL,
                },
                rpc.outbox[0].node_id
            )
        )
        await asyncio.sleep(1)
        assert worker.user_count == 2
        # master send
        await rpc.mocked_send(Message("stop", None, rpc.outbox[0].node_id))
        await asyncio.sleep(1)
        assert "stopped" in [msg.type for msg in rpc.outbox]
        assert worker.user_count == 0
        assert worker.worker_state == STATE_INIT
        # master send
        await rpc.mocked_send(Message("quit", None, rpc.outbox[0].node_id))
        await asyncio.sleep(4)
        assert len(worker.tasks) == 1
        assert worker.tasks[0].get_name() == "worker_run"

aiotest/test/test_runner.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import traceback
import pytest
import asyncio
import allure
from aiotest.user import AsyncHttpUser
from aiotest.stats_exporter import parse_error
from aiotest import events
from aiotest.test.setting import BASE_URL, P1, P2, P3
from aiotest import runners


request_dict = {}
error_dict = {}


async def on_request(runner, request_name, request_method, response_time, response_length, error):
    global request_dict
    request_dict = {}
    if error:
        error = parse_error(error)
    request_dict = {
        "request_name": request_name, 
        "request_method": request_method, 
        "response_time": response_time,
        "response_length": response_length,
        "error": error,
    }


async def on_error(runner, error):
    global error_dict
    error_dict = {}
    if error:
        error = parse_error(error)
    error_dict = {
        "error": error,
    }


@allure.title("add/remove events.stats_request/user_error")
@pytest.fixture(scope="module", autouse=True)
def events_handlers():
    events.stats_request += on_request
    events.user_error += on_error
    yield
    events.stats_request -= on_request
    events.user_error -= on_error

@allure.title("init request_dict/error_dict")
@pytest.fixture(scope="function", autouse=True)
def init_request_and_error_dict():
    global request_dict, error_dict
    request_dict = {}
    error_dict = {}



@allure.feature("aiotest/user")
@allure.story("aiotest/user/AsyncHttpUser")
class TestAsyncHttpUser:
        
    @allure.severity(P1)
    @allure.title("user jobs")
    async def test_01(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                pass
            async def get_test(self):
                pass
            async def get(self):
                pass
        u =MyUser()
        assert u.jobs == [MyUser.test_get, MyUser.get_test]
        u.on_stop()

    @allure.severity(P2)
    @allure.title("on_start add job")
    async def test_02(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def on_start(self):
                self.jobs = [MyUser.get01]
            async def get01(self):
                async with self.session.get("/") as resp: 
                    text = await resp.text()
            async def get02(self):
                pass
            async def get03(self):
                pass
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert u.jobs == [MyUser.get01]
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 150 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"] == None
        assert error_dict.get("error", None) is None
                
    @allure.severity(P1)
    @allure.title("user start jobs")
    async def test_03(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp: 
                    text = await resp.text()
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 120 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"] == None
        assert error_dict.get("error", None) is None

    @allure.severity(P2)
    @allure.title("user on_start raise error")
    async def test_04(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def on_start(self):
                async with self.session.get("/") as resp: 
                    text = await resp.text()
                    raise KeyError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 130 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"].find("KeyError") >= 0
        assert error_dict.get("error", None) is None

    @allure.severity(P1)
    @allure.title("user.job raise error")
    async def test_05(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp: 
                    text = await resp.text()
                    raise KeyError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 120 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"].find("KeyError") >= 0
        assert error_dict.get("error", None) is None

    @allure.severity(P2)
    @allure.title("user.job raise NameError error")
    async def test_06(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/fast") as resp:
                    text = await resp.text() 
                    raise NameError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict == {}
        assert error_dict.get("error", None).find("NameError") >= 0
        u.on_stop()

    @allure.severity(P2)
    @allure.title("user.job raise KeyError error")
    async def test_07(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp:
                    text = await resp.text() 
                    raise KeyError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 130 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"].find("KeyError") >= 0
        assert error_dict.get("error", None) is None

    @allure.severity(P2)
    @allure.title("user.job raise CancelledError error")
    async def test_08(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp:
                    text = await resp.text() 
                    raise asyncio.CancelledError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        with pytest.raises(asyncio.exceptions.CancelledError):
            await u.start()
        assert request_dict == {}
        assert error_dict.get("error", None) is None

    @allure.severity(P2)
    @allure.title("user.job success() and raise KeyError error")
    async def test_09(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp:
                    text = await resp.text() 
                    resp.success()
                    raise KeyError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 120 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"] is None
        assert error_dict.get("error", None).find("KeyError") >= 0
        u.on_stop()

    @allure.severity(P2)
    @allure.title("user.job failure() and raise KeyError error")
    async def test_10(self):
        class MyUser(AsyncHttpUser):
            host = BASE_URL
            wait_time = 0
            async def test_get(self):
                async with self.session.get("/") as resp:
                    text = await resp.text() 
                    resp.failure("my failuer")
                    raise KeyError
            async def start(self):
                try:
                    await self.on_start()
                    await asyncio.sleep(self.wait_time)
                    for job in self.jobs:
                        await job(self)
                        await asyncio.sleep(self.wait_time)
                except asyncio.CancelledError:
                    await self.on_stop()
                    raise
                except Exception as e:
                    await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())

        u =MyUser()
        await u.start()
        assert request_dict["request_name"] == "/"
        assert request_dict["request_method"] == "GET"
        assert 120 > request_dict["response_time"] > 80
        assert request_dict["response_length"] == len("asyncio sleep")  
        assert request_dict["error"].find("my failuer") >= 0
        assert error_dict.get("error", None).find("KeyError") >= 0
        u.on_stop()












aiotest/test/test_user.py
hewei198711/Aiotest
# encoding: utf-8

import sys
import asyncio
import traceback

from aiohttp.client_exceptions import ClientResponseError
from aiotest import events
from aiotest.client import ClientSession
from aiotest import runners


class UserMeta(type):
    """
    Meta class for the main User class. It's used to allow User classes to specify task execution 
    """
    def __new__(mcs, classname, bases, class_dict):
        jobs = []
        for key, value in class_dict.items():
            if key.startswith("test") or key.endswith("test"):
                jobs.append(value)
        class_dict["jobs"] = jobs
        
        return type.__new__(mcs, classname, bases, class_dict)


class User(metaclass=UserMeta):
    host = None
    session = None
    wait_time = 1
    weight = 1
    task = None
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    async def on_start(self):
        """
        Called when a User starts running.
        """
        pass
    
    async def on_stop(self):
        """
        Called when a User stops running (is canceled)
        """
        if not self.session.closed:
            await self.session.close()
        self.session = None
    
    def start_user(self):
        self.task = asyncio.create_task(self.start(), name=type(self).__name__)
        
    async def stop_user(self):
        if not self.task.done():
            self.task.cancel()
        await self.on_stop()
                
    async def start(self):
        """
        Synchronization executes coroutines job from top to bottom
        """
        try:
            try:
                await self.on_start()
                await asyncio.sleep(self.wait_time)
            except:
                await self.on_stop()
                raise                
            while True:
                for job in self.jobs:
                    await job(self)
                    await asyncio.sleep(self.wait_time)
        except asyncio.CancelledError:
            await self.on_stop()
            raise
        except Exception as e:
            await events.user_error.fire(runner=runners.global_runner, error=f"{sys.exc_info()[0].__name__}: {e}" + "".join(traceback.format_tb(sys.exc_info()[2])).strip())


class AsyncHttpUser(User):
    """
    Represents an AsyncHttp "user" which is to be started and attack the system that is to be load tested.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.session = ClientSession(base_url=self.host)

aiotest/user.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser, events

"""
events:
    init_command_line_parser
    init
    test_start
    start_complete
    stats_request
    user_error
    worker_report
    test_stop
    quitting

command: aiotest -f add_events.py -u 2 -r 2 -t 120
View test data: http://localhost:8089
"""

async def on_init_command_line_parser(parser):
    parser.add_argument(
        '--my-order',
        type=str,
        default='',
        help="It is my order"
    )

async def on_init():
    print(f"init...")

async def on_test_start(runner):
    print("test start...")

async def on_start_complete(user_count, runner):
    print(f"Number of simulated users: {user_count}")

async def on_stats_request(runner, request_name, request_method, response_time, response_length, error):
    stats = {
        "request_name": request_name, 
        "request_method": request_method, 
        "response_time": response_time,
        "response_length": response_length,
        "error": error,
        "user_count": runner.user_count,
    }
    print(stats)

async def on_user_error(runner, error):
    print(error)

async def on_worker_report(runner, worker_id, data):
    print(f"{worker_id}: {data}")

async def on_test_stop(runner):
    print("test stop...")

async def on_quitting():
    print("quitting...")


events.init_command_line_parser += on_init_command_line_parser
events.init += on_init
events.test_start += on_test_start
events.start_complete += on_start_complete
events.stats_request += on_stats_request
events.user_error += on_user_error
events.worker_report += on_worker_report
events.test_stop += on_test_stop
events.quitting += on_quitting


class TestUser(AsyncHttpUser):
    host = "https://httpbin.org"

    async def test_post(self):
        url="/post"
        async with self.session.post(url=url) as resp:
            data = await resp.json()
    
    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()
   


examples/add_events.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser, logger

# command: aiotest -f add_logger.py -u 2 -r 2 -t 120 -L DEBUG --logfile log.log
# View test data: http://localhost:8089 

class TestUser(AsyncHttpUser):
    wait_time = 1
    weight = 1
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()
            logger.debug(self.token)

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()
            logger.debug(data)

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()
            logger.debug(data)


examples/add_logger.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser

# command: aiotest -f basic.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 

class TestUser(AsyncHttpUser):
    wait_time = 1
    weight = 1
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()


examples/basic.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser

# command: aiotest -f custom_wait_time.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 

class TestUser(AsyncHttpUser):
    # The default wait is 1 second
    wait_time = 0 # not wait
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()



examples/custom_wait_time.py
hewei198711/Aiotest
# encoding: utf-8

from redis import StrictRedis
from aiotest import AsyncHttpUser

# command: aiotest -f distributed_test.py --master --expect-workers 4 --master-bind-host "127.0.0.1" --master-bind-port 5557 -u 400 -r 80 -t 3600
# command: aiotest -f distributed_test.py --worker --master-host "127.0.0.1" -master-port 5557
# command: aiotest -f distributed_test.py --worker --master-host "127.0.0.1" -master-port 5557
# command: aiotest -f distributed_test.py --worker --master-host "127.0.0.1" -master-port 5557
# command: aiotest -f distributed_test.py --worker --master-host "127.0.0.1" -master-port 5557

# The data is placed in redis, which is very useful in distributed testing, 
# and can easily ensure the load balance of various working nodes
# View test data: http://localhost:8089 

db = StrictRedis(
    host="127.0.0.1", 
    port=6379, 
    db=0, 
    decode_responses=True, 
    password="123456"
)
db.pipeline()

class TestUser(AsyncHttpUser):
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        username = self.db.lpop("user")

        url="/login"
        data = {"username": username, "password": 123456}
        async with self.session.post(url=url, data=data) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

examples/distributed_test.py
hewei198711/Aiotest
# encoding: utf-8

import sys
from aiotest import AsyncHttpUser
from redis import StrictRedis

# The data is placed in redis, which is very useful in distributed testing, 
# and can easily ensure the load balance of various working nodes

# command: aiotest -f get_data_from_redis.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 


db = StrictRedis(
    host="127.0.0.1", 
    port=6379, 
    db=0, 
    decode_responses=True, 
    password="123456"
)
db.pipeline()

class TestUser(AsyncHttpUser):
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        username = self.db.lpop("user")

        url="/login"
        data = {"username": username, "password": 123456}
        async with self.session.post(url=url, data=data) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()



examples/get_data_from_redis.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser

# The execution weight of TestUser01 : TestUser02 is 2:1
# command: aiotest -f multiple_user_classes.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 

class TestUser01(AsyncHttpUser):
    # The default weight is 1 
    weight = 2
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()


class TestUser02(AsyncHttpUser):
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()
examples/multiple_user_classes.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser

# command: aiotest -f name_property.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 

class TestUser(AsyncHttpUser):
    wait_time = 1
    weight = 1
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url, name="custom_name_one") as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url, name="custom_name_two") as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url, name="custom_name_three") as resp:
            data = await resp.text()


examples/name_property.py
hewei198711/Aiotest
# encoding: utf-8

from asyncio.locks import Semaphore
from aiotest import AsyncHttpUser, events

# Second kill activity, all users have logged in, complete the collection, and then synchronize to perform subsequent tests.
# command: aiotest -f semaphore.py -u 20 -r 2 -t 1200
# View test data: http://localhost:8089 


all_user_start_complete = Semaphore()
all_user_start_complete.acquire()


async def on_start_complete(user_count, runner):
    all_user_start_complete.release()

          
events.start_complete += on_start_complete 
          
    
class TestUser(AsyncHttpUser):
    wait_time = 1
    weight = 1
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()
        all_user_start_complete.wait()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()


examples/semaphore.py
hewei198711/Aiotest
# encoding: utf-8

from aiotest import AsyncHttpUser, LoadUserShape

# command: aiotest -f stepwise_pressure_shape.py -u 2 -r 2 -t 120
# View test data: http://localhost:8089 

class TestUser(AsyncHttpUser):
    host = "https://httpbin.org"
    token = None

    async def on_start(self):
        "login and get the token"
        url="/post"
        async with self.session.post(url=url) as resp:
            self.token = await resp.json()

    async def test_get(self):
        url = "/get"
        async with self.session.get(url=url) as resp:
            data = await resp.text()

    async def test_post(self):
        url = "/status/200"
        async with self.session.post(url=url) as resp:
            data = await resp.text()


                   
class TestShape(LoadUserShape):
    stages = [
        {"duration": 600, "user_count": 1000, "rate": 50},
        {"duration": 1200, "user_count": 3000, "rate": 50},
        {"duration": 1800, "user_count": 5000, "rate": 50},
        {"duration": 2400, "user_count": 3000, "rate": 50},
        {"duration": 3600, "user_count": 1000, "rate": 50},
    ]

    def tick(self):
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                tick_data = (stage["user_count"], stage["rate"])
                return tick_data

        return None
examples/stepwise_pressure_shape.py
hewei198711/Aiotest
# encoding: utf-8

import pytest

if __name__ == "__main__":
    pytest.main([ 
        "-v",
        "--asyncio-mode",
        "auto",
        "--alluredir",
        "aiotest/test/reports",
        "aiotest/test",
        "--clean-alluredir",
        "--disable-warnings",
        "--color",
        "yes"
    ])

pytest_run.py
hewei198711/Aiotest