Fix bug in rlib2 unit tests

[ganeti-local] / lib / rapi / testutils.py

diff --git a/lib/rapi/testutils.py b/lib/rapi/testutils.py
index f960381..cdc34dc 100644 (file)
--- a/lib/rapi/testutils.py
+++ b/lib/rapi/testutils.py
@@ -25,7 +25,6 @@
 
 import logging
 import re
-import mimetools
 import base64
 import pycurl
 from cStringIO import StringIO
@@ -135,6 +134,17 @@ def _GetPathFromUri(uri):
     return None
 
 
+def _FormatHeaders(headers):
+  """Formats HTTP headers.
+
+  @type headers: sequence of strings
+  @rtype: string
+
+  """
+  assert compat.all(": " in header for header in headers)
+  return "\n".join(headers)
+
+
 class FakeCurl:
   """Fake cURL object.
 
@@ -168,11 +178,11 @@ class FakeCurl:
     writefn = self._opts[pycurl.WRITEFUNCTION]
 
     if pycurl.HTTPHEADER in self._opts:
-      baseheaders = "\n".join(self._opts[pycurl.HTTPHEADER])
+      baseheaders = _FormatHeaders(self._opts[pycurl.HTTPHEADER])
     else:
       baseheaders = ""
 
-    headers = mimetools.Message(StringIO(baseheaders), 0)
+    headers = http.ParseHeaders(StringIO(baseheaders))
 
     if request_body:
       headers[http.HTTP_CONTENT_LENGTH] = str(len(request_body))
@@ -188,7 +198,7 @@ class FakeCurl:
         "%s %s" % (http.auth.HTTP_BASIC_AUTH, base64.b64encode(userpwd))
 
     path = _GetPathFromUri(url)
-    (code, resp_body) = \
+    (code, _, resp_body) = \
       self._handler.FetchResponse(path, method, headers, request_body)
 
     self._info[pycurl.RESPONSE_CODE] = code
@@ -200,7 +210,7 @@ class _RapiMock:
   """Mocking out the RAPI server parts.
 
   """
-  def __init__(self, user_fn, luxi_client):
+  def __init__(self, user_fn, luxi_client, reqauth=False):
     """Initialize this class.
 
     @type user_fn: callable
@@ -209,7 +219,7 @@ class _RapiMock:
 
     """
     self.handler = \
-      server.rapi.RemoteApiHandler(user_fn, _client_cls=luxi_client)
+      server.rapi.RemoteApiHandler(user_fn, reqauth, _client_cls=luxi_client)
 
   def FetchResponse(self, path, method, headers, request_body):
     """This is a callback method used to fetch a response.
@@ -224,7 +234,7 @@ class _RapiMock:
     @param request_body: Request body
     @type headers: mimetools.Message
     @param headers: Request headers
-    @return: Tuple containing status code and response body
+    @return: Tuple containing status code, response headers and response body
 
     """
     req_msg = http.HttpMessage()
@@ -236,7 +246,7 @@ class _RapiMock:
     (_, _, _, resp_msg) = \
       http.server.HttpResponder(self.handler)(lambda: (req_msg, None))
 
-    return (resp_msg.start_line.code, resp_msg.body)
+    return (resp_msg.start_line.code, resp_msg.headers, resp_msg.body)
 
 
 class _TestLuxiTransport:
@@ -294,7 +304,7 @@ class _LuxiCallRecorder:
     """
     return self._called
 
-  def __call__(self):
+  def __call__(self, address=None):
     """Creates an instrumented LUXI client.
 
     The LUXI client will record all method calls (use L{CalledNames} to
@@ -302,7 +312,8 @@ class _LuxiCallRecorder:
 
     """
     return luxi.Client(transport=compat.partial(_TestLuxiTransport,
-                                                self.Record))
+                                                self.Record),
+                       address=address)
 
 
 def _TestWrapper(fn, *args, **kwargs):