mercurial/httpclient/tests/util.py
changeset 16775 e6af8676302f
parent 16774 69af967b6d6f
child 16776 5088d0b9a9a1
equal deleted inserted replaced
16774:69af967b6d6f 16775:e6af8676302f
     1 # Copyright 2010, Google Inc.
       
     2 # All rights reserved.
       
     3 #
       
     4 # Redistribution and use in source and binary forms, with or without
       
     5 # modification, are permitted provided that the following conditions are
       
     6 # met:
       
     7 #
       
     8 #     * Redistributions of source code must retain the above copyright
       
     9 # notice, this list of conditions and the following disclaimer.
       
    10 #     * Redistributions in binary form must reproduce the above
       
    11 # copyright notice, this list of conditions and the following disclaimer
       
    12 # in the documentation and/or other materials provided with the
       
    13 # distribution.
       
    14 #     * Neither the name of Google Inc. nor the names of its
       
    15 # contributors may be used to endorse or promote products derived from
       
    16 # this software without specific prior written permission.
       
    17 
       
    18 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
       
    19 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
       
    20 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
       
    21 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
       
    22 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
       
    23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
       
    24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
       
    25 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
       
    26 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
       
    27 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
       
    28 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
       
    29 import difflib
       
    30 import socket
       
    31 
       
    32 import httpplus
       
    33 
       
    34 
       
    35 class MockSocket(object):
       
    36     """Mock non-blocking socket object.
       
    37 
       
    38     This is ONLY capable of mocking a nonblocking socket.
       
    39 
       
    40     Attributes:
       
    41       early_data: data to always send as soon as end of headers is seen
       
    42       data: a list of strings to return on recv(), with the
       
    43             assumption that the socket would block between each
       
    44             string in the list.
       
    45       read_wait_sentinel: data that must be written to the socket before
       
    46                           beginning the response.
       
    47       close_on_empty: If true, close the socket when it runs out of data
       
    48                       for the client.
       
    49     """
       
    50     def __init__(self, af, socktype, proto):
       
    51         self.af = af
       
    52         self.socktype = socktype
       
    53         self.proto = proto
       
    54 
       
    55         self.early_data = []
       
    56         self.data = []
       
    57         self.remote_closed = self.closed = False
       
    58         self.close_on_empty = False
       
    59         self.sent = ''
       
    60         self.read_wait_sentinel = httpplus._END_HEADERS
       
    61         self.blocking = True
       
    62 
       
    63     def close(self):
       
    64         self.closed = True
       
    65 
       
    66     def connect(self, sa):
       
    67         self.sa = sa
       
    68 
       
    69     def setblocking(self, timeout):
       
    70         self.blocking = bool(timeout)
       
    71 
       
    72     def recv(self, amt=-1):
       
    73         # we only properly emulate non-blocking sockets
       
    74         assert not self.blocking
       
    75         if self.early_data:
       
    76             datalist = self.early_data
       
    77         elif not self.data:
       
    78             return ''
       
    79         else:
       
    80             datalist = self.data
       
    81         if amt == -1:
       
    82             return datalist.pop(0)
       
    83         data = datalist.pop(0)
       
    84         if len(data) > amt:
       
    85             datalist.insert(0, data[amt:])
       
    86         if not self.data and not self.early_data and self.close_on_empty:
       
    87             self.remote_closed = True
       
    88         return data[:amt]
       
    89 
       
    90     @property
       
    91     def ready_for_read(self):
       
    92         return ((self.early_data and httpplus._END_HEADERS in self.sent)
       
    93                 or (self.read_wait_sentinel in self.sent and self.data)
       
    94                 or self.closed or self.remote_closed)
       
    95 
       
    96     def send(self, data):
       
    97         # this is a horrible mock, but nothing needs us to raise the
       
    98         # correct exception yet
       
    99         assert not self.closed, 'attempted to write to a closed socket'
       
   100         assert not self.remote_closed, ('attempted to write to a'
       
   101                                         ' socket closed by the server')
       
   102         if len(data) > 8192:
       
   103             data = data[:8192]
       
   104         self.sent += data
       
   105         return len(data)
       
   106 
       
   107 
       
   108 def mockselect(r, w, x, timeout=0):
       
   109     """Simple mock for select()
       
   110     """
       
   111     readable = filter(lambda s: s.ready_for_read, r)
       
   112     return readable, w[:], []
       
   113 
       
   114 
       
   115 class MockSSLSocket(object):
       
   116     def __init__(self, sock):
       
   117         self._sock = sock
       
   118         self._fail_recv = True
       
   119 
       
   120     def __getattr__(self, key):
       
   121         return getattr(self._sock, key)
       
   122 
       
   123     def __setattr__(self, key, value):
       
   124         if key not in ('_sock', '_fail_recv'):
       
   125             return setattr(self._sock, key, value)
       
   126         return object.__setattr__(self, key, value)
       
   127 
       
   128     def recv(self, amt=-1):
       
   129         try:
       
   130             if self._fail_recv:
       
   131                 raise socket.sslerror(socket.SSL_ERROR_WANT_READ)
       
   132             return self._sock.recv(amt=amt)
       
   133         finally:
       
   134             self._fail_recv = not self._fail_recv
       
   135 
       
   136 
       
   137 def mocksslwrap(sock, keyfile=None, certfile=None,
       
   138                 server_side=False, cert_reqs=httpplus.socketutil.CERT_NONE,
       
   139                 ssl_version=None, ca_certs=None,
       
   140                 do_handshake_on_connect=True,
       
   141                 suppress_ragged_eofs=True):
       
   142     assert sock.blocking, ('wrapping a socket with ssl requires that '
       
   143                            'it be in blocking mode.')
       
   144     return MockSSLSocket(sock)
       
   145 
       
   146 
       
   147 def mockgetaddrinfo(host, port, unused, streamtype):
       
   148     assert unused == 0
       
   149     assert streamtype == socket.SOCK_STREAM
       
   150     if host.count('.') != 3:
       
   151         host = '127.0.0.42'
       
   152     return [(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, '',
       
   153              (host, port))]
       
   154 
       
   155 
       
   156 class HttpTestBase(object):
       
   157     def setUp(self):
       
   158         self.orig_socket = socket.socket
       
   159         socket.socket = MockSocket
       
   160 
       
   161         self.orig_getaddrinfo = socket.getaddrinfo
       
   162         socket.getaddrinfo = mockgetaddrinfo
       
   163 
       
   164         self.orig_select = httpplus.select.select
       
   165         httpplus.select.select = mockselect
       
   166 
       
   167         self.orig_sslwrap = httpplus.socketutil.wrap_socket
       
   168         httpplus.socketutil.wrap_socket = mocksslwrap
       
   169 
       
   170     def tearDown(self):
       
   171         socket.socket = self.orig_socket
       
   172         httpplus.select.select = self.orig_select
       
   173         httpplus.socketutil.wrap_socket = self.orig_sslwrap
       
   174         socket.getaddrinfo = self.orig_getaddrinfo
       
   175 
       
   176     def assertStringEqual(self, l, r):
       
   177         try:
       
   178             self.assertEqual(l, r, ('failed string equality check, '
       
   179                                     'see stdout for details'))
       
   180         except:
       
   181             add_nl = lambda li: map(lambda x: x + '\n', li)
       
   182             print 'failed expectation:'
       
   183             print ''.join(difflib.unified_diff(
       
   184                 add_nl(l.splitlines()), add_nl(r.splitlines()),
       
   185                 fromfile='expected', tofile='got'))
       
   186             raise
       
   187 
       
   188     def doPost(self, con, expect_body, body_to_send='This is some POST data'):
       
   189         con.request('POST', '/', body=body_to_send,
       
   190                     expect_continue=True)
       
   191         expected_req = ('POST / HTTP/1.1\r\n'
       
   192                         'Host: 1.2.3.4\r\n'
       
   193                         'content-length: %d\r\n'
       
   194                         'Expect: 100-Continue\r\n'
       
   195                         'accept-encoding: identity\r\n\r\n' %
       
   196                         len(body_to_send))
       
   197         if expect_body:
       
   198             expected_req += body_to_send
       
   199         return expected_req
       
   200 # no-check-code