diff --git a/doc/newsfragments/2597_changed_tcpclient_interface.rst b/doc/newsfragments/2597_changed_tcpclient_interface.rst new file mode 100755 index 000000000..095b3784d --- /dev/null +++ b/doc/newsfragments/2597_changed_tcpclient_interface.rst @@ -0,0 +1 @@ +TCPClient supports binding interface \ No newline at end of file diff --git a/testplan/common/utils/sockets/client.py b/testplan/common/utils/sockets/client.py index c4ebd186d..6674584c4 100644 --- a/testplan/common/utils/sockets/client.py +++ b/testplan/common/utils/sockets/client.py @@ -2,6 +2,7 @@ import time import socket +from typing import Union, Tuple class Client: @@ -16,19 +17,21 @@ class Client: 4. close """ - def __init__(self, host, port, interface=None): + def __init__( + self, + host: str, + port: Union[str, int], + interface: Union[Tuple[str, int], None] = None, + ) -> None: """ Create a new TCP client. This constructor takes parameters that specify the address (host, port) to connect to and an optional logging callback method. :param host: hostname or IP address to connect to - :type host: ``str`` :param port: port to connect to - :type port: ``str`` or ``int`` :param interface: Local interface to bind to. Defaults to None, in which case the socket does not bind before connecting. - :type interface: (``str``, ``str`` or ``int``) tuple """ self._input_host = host self._input_port = port @@ -37,49 +40,44 @@ def __init__(self, host, port, interface=None): self._timeout = None @property - def address(self): + def address(self) -> Tuple[str, int]: """ Returns the host and port information of socket. """ return self._client.getsockname() @property - def port(self): + def port(self) -> Union[str, int]: return self._input_port - def connect(self): + def connect(self) -> None: """ Connect client to socket. """ self._client = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - if self._interface is not None: + if self._interface: self._client.bind(self._interface) self._client.connect((self._input_host, self._input_port)) - def send(self, msg): + def send(self, msg: bytes) -> Tuple[float, int]: """ Send the given message. :param msg: Message to be sent. - :type msg: ``bytes`` :return: Timestamp when msg sent (in microseconds from epoch) and number of bytes sent - :rtype: ``tuple`` of ``long`` and ``int`` """ tsp = time.time() * 1000000 size = self._client.send(msg) return tsp, size - def receive(self, size, timeout=30): + def receive(self, size: int, timeout: int = 30) -> bytes: """ Receive a message. :param size: Number of bytes to receive. - :type size: ``int`` :param timeout: Timeout in seconds. - :type timeout: ``int`` :return: message received - :rtype: ``bytes`` """ if timeout != self._timeout: self._timeout = timeout @@ -92,25 +90,19 @@ def receive(self, size, timeout=30): raise return msg - def recv(self, bufsize, flags=0): + def recv(self, bufsize: int, flags: int = 0) -> bytes: """ Proxy for Python's ``socket.recv()``. :param bufsize: Maximum amount of data to be received at once. - :type bufsize: ``int`` :param flags: Defaults to zero. - :type flags: ``int`` :return: message received - :rtype: ``bytes`` """ return self._client.recv(bufsize, flags) - def close(self): + def close(self) -> None: """ Close the connection. - - :return: ``None`` - :rtype: ``NoneType`` """ if self._client is not None: self._client.close() diff --git a/testplan/testing/multitest/driver/tcp/client.py b/testplan/testing/multitest/driver/tcp/client.py index 3bb7efe27..aabe17c0e 100644 --- a/testplan/testing/multitest/driver/tcp/client.py +++ b/testplan/testing/multitest/driver/tcp/client.py @@ -43,19 +43,14 @@ class TCPClient(Driver): {emphasized_members_docs} :param name: Name of TCPClient. - :type name: ``str`` :param host: Target host name. This can be a :py:class:`~testplan.common.utils.context.ContextValue` and will be expanded on runtime. - :type host: ``str`` :param port: Target port number. This can be a :py:class:`~testplan.common.utils.context.ContextValue` and will be expanded on runtime. - :type port: ``int`` :param interface: Interface to bind to. - :type interface: ``NoneType`` or ``tuple``(``str, ``int``) :param connect_at_start: Connect to server on start. Default: True - :type connect_at_start: ``bool`` Also inherits all :py:class:`~testplan.testing.multitest.driver.base.Driver` options. @@ -68,73 +63,69 @@ def __init__( name: str, host: Union[str, ContextValue], port: Union[int, str, ContextValue], - interface: Optional[Union[str, Tuple[str, int]]] = None, + interface: Union[Tuple[str, int], None] = None, connect_at_start: bool = True, **options ): options.update(self.filter_locals(locals())) super(TCPClient, self).__init__(**options) - self._host: str = None - self._port: int = None + self._host: Optional[str] = None + self._port: Optional[int] = None self._client = None self._server_host = None self._server_port = None @property - def host(self): + def host(self) -> str: """Target host name.""" return self._host @property - def port(self): + def port(self) -> int: """Client port number assigned.""" return self._port @property - def server_port(self): + def server_port(self) -> int: return self._server_port - def connect(self): + def connect(self) -> None: """ Connect client. """ self._client.connect() self._host, self._port = self._client.address - def send_text(self, msg, standard="utf-8"): + def send_text(self, msg: str, standard: str = "utf-8") -> int: """ Encodes to bytes and calls :py:meth:`TCPClient.send `. """ - return self.send(bytes(msg.encode(standard))) + return self.send(msg.encode(standard)) - def send(self, msg): + def send(self, msg: bytes) -> int: """ Sends bytes. :param msg: Message to be sent - :type msg: ``bytes`` :return: Number of bytes sent - :rtype: ``int`` """ return self._client.send(msg)[1] - def send_tsp(self, msg): + def send_tsp(self, msg: bytes) -> Tuple[float, int]: """ Sends bytes and returns also timestamp sent. :param msg: Message to be sent - :type msg: ``bytes`` :return: Timestamp when msg sent (in microseconds from epoch) and number of bytes sent - :rtype: ``tuple`` of ``long`` and ``int`` """ return self._client.send(msg) - def receive_text(self, standard="utf-8", **kwargs): + def receive_text(self, standard: str = "utf-8", **kwargs) -> str: """ Calls :py:meth:`TCPClient.receive @@ -143,7 +134,7 @@ def receive_text(self, standard="utf-8", **kwargs): """ return self.receive(**kwargs).decode(standard) - def receive(self, size=1024, timeout=30): + def receive(self, size: int = 1024, timeout: int = 30) -> Optional[bytes]: """Receive bytes from the given connection.""" received = None timeout_info = TimeoutExceptionInfo() @@ -158,34 +149,38 @@ def receive(self, size=1024, timeout=30): ) return received - def reconnect(self): + def reconnect(self) -> None: """Client reconnect.""" self.close() self.connect() - def starting(self): + def starting(self) -> None: """Start the TCP client and optionally connect to host/post.""" super(TCPClient, self).starting() self._server_host = expand(self.cfg.host, self.context) self._server_port = networking.port_to_int( expand(self.cfg.port, self.context) ) - self._client = Client(host=self._server_host, port=self._server_port) + self._client = Client( + host=self._server_host, + port=self._server_port, + interface=self.cfg.interface, + ) if self.cfg.connect_at_start: self.connect() - def stopping(self): + def stopping(self) -> None: """Close the client connection.""" super(TCPClient, self).stopping() self.close() - def close(self): + def close(self) -> None: """ Close connection. """ if self._client: self._client.close() - def aborting(self): + def aborting(self) -> None: """Abort logic that stops the client.""" self.close()