mercredi 27 avril 2016

How to mock a socket object via the mock library

How to mock a TCPSocket wrapper for the socket from the Python's standard libary via the mock library (unittest.mock in case of Python 3)?

This is my wrapper:

import socket

import utils


class TCPSocket:
    def __init__(self):
        self.buf = ''

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self, host, port):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((host, port))

    def close(self):
        self.sock.close()

    def send(self, data):
        self.sock.send(data)

    def sendall(self, data):
        self.sock.sendall(data)

    # For best match with hardware and network realities,
    # the value of limit should be a relatively small power of 2, for example, 4096
    def recv_some(self, limit=4096):
        return self.sock.recv(limit)

    def recv_bytes(self, count):
        while len(self.buf) < count:
            self.buf += self.recv_some()
        parts = utils.split_first_n(self.buf, count)
        self.buf = parts[1]
        return parts[0]

    def recv_until(self, delim):
        while delim not in self.buf:
            self.buf += self.recv_some()
        parts = self.buf.split(delim, maxsplit=1)
        self.buf = parts[1]
        return parts[0]

I want to test whether functions like recv_until and recv_bytes do what they really need.

Aucun commentaire:

Enregistrer un commentaire