Synnefo » snf-ganeti

class TestSocketWrapper(object):

  """ Utility class that opens a domain socket and cleans up as needed.

  """

  def __init__(self):

    """ Constructor cleaning up variables to be used.

    """

    self.tmpdir = None

    self.sock = None

  def Create(self, max_connections=1):

    """ Creates a bound and ready socket, cleaning up in case of failure.

    @type max_connections: int

    @param max_connections: The number of max connections allowed for the

                            socket.

    @rtype: tuple of socket, string

    @return: The socket object and the path to reach it with.

    """

    # Using a temporary directory as there's no easy way to create temporary

    # sockets without writing a custom loop around tempfile.mktemp and

    # socket.bind

    self.tmpdir = tempfile.mkdtemp()

    try:

      tmpsock = utils.PathJoin(self.tmpdir, "sock")

      logging.debug("Creating temporary socket at %s", tmpsock)

      self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

      try:

        self.sock.bind(tmpsock)

        self.sock.listen(max_connections)

      except:

        self.sock.close()

        raise

    except:

      shutil.rmtree(self.tmpdir)

      raise

    return self.sock, tmpsock

  def Destroy(self):

    """ Destroys the socket and performs all necessary cleanup.

    """

    if self.tmpdir is None or self.sock is None:

      raise Exception("A socket must be created successfully before attempting "

                      "its destruction")

    try:

      self.sock.close()

    finally:

      shutil.rmtree(self.tmpdir)

    socket_wrapper = TestSocketWrapper()

    sock, path = socket_wrapper.Create()

    cb(path)

    try:

      sock.settimeout(cls._CLIENT_CONNECT_TIMEOUT)

      (conn, _) = sock.accept()

    except socket.error, err:

      raise errcls("Client didn't connect in time (%s)" % err)

      socket_wrapper.Destroy()