/atom/mock_http_core.py

  1#!/usr/bin/env python
  2#
  3# Copyright (C) 2009 Google Inc.
  4#
  5# Licensed under the Apache License, Version 2.0 (the "License");
  6# you may not use this file except in compliance with the License.
  7# You may obtain a copy of the License at
  8#
  9#      http://www.apache.org/licenses/LICENSE-2.0
 10#
 11# Unless required by applicable law or agreed to in writing, software
 12# distributed under the License is distributed on an "AS IS" BASIS,
 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 14# See the License for the specific language governing permissions and
 15# limitations under the License.
 16
 17
 18# This module is used for version 2 of the Google Data APIs.
 19
 20
 21__author__ = 'j.s@google.com (Jeff Scudder)'
 22
 23
 24import StringIO
 25import pickle
 26import os.path
 27import tempfile
 28import atom.http_core
 29
 30
 31class Error(Exception):
 32  pass
 33
 34
 35class NoRecordingFound(Error):
 36  pass
 37
 38
 39class MockHttpClient(object):
 40  debug = None
 41  real_client = None
 42  last_request_was_live = False
 43
 44  # The following members are used to construct the session cache temp file
 45  # name.
 46  # These are combined to form the file name
 47  # /tmp/cache_prefix.cache_case_name.cache_test_name
 48  cache_name_prefix = 'gdata_live_test'
 49  cache_case_name = ''
 50  cache_test_name = ''
 51
 52  def __init__(self, recordings=None, real_client=None):
 53    self._recordings = recordings or []
 54    if real_client is not None:
 55      self.real_client = real_client
 56
 57  def add_response(self, http_request, status, reason, headers=None,
 58      body=None):
 59    response = MockHttpResponse(status, reason, headers, body)
 60    # TODO Scrub the request and the response.
 61    self._recordings.append((http_request._copy(), response))
 62
 63  AddResponse = add_response
 64
 65  def request(self, http_request):
 66    """Provide a recorded response, or record a response for replay.
 67
 68    If the real_client is set, the request will be made using the
 69    real_client, and the response from the server will be recorded.
 70    If the real_client is None (the default), this method will examine
 71    the recordings and find the first which matches.
 72    """
 73    request = http_request._copy()
 74    _scrub_request(request)
 75    if self.real_client is None:
 76      self.last_request_was_live = False
 77      for recording in self._recordings:
 78        if _match_request(recording[0], request):
 79          return recording[1]
 80    else:
 81      # Pass along the debug settings to the real client.
 82      self.real_client.debug = self.debug
 83      # Make an actual request since we can use the real HTTP client.
 84      self.last_request_was_live = True
 85      response = self.real_client.request(http_request)
 86      scrubbed_response = _scrub_response(response)
 87      self.add_response(request, scrubbed_response.status,
 88                        scrubbed_response.reason,
 89                        dict(atom.http_core.get_headers(scrubbed_response)),
 90                        scrubbed_response.read())
 91      # Return the recording which we just added.
 92      return self._recordings[-1][1]
 93    raise NoRecordingFound('No recoding was found for request: %s %s' % (
 94        request.method, str(request.uri)))
 95
 96  Request = request
 97
 98  def _save_recordings(self, filename):
 99    recording_file = open(os.path.join(tempfile.gettempdir(), filename),
100                          'wb')
101    pickle.dump(self._recordings, recording_file)
102    recording_file.close()
103
104  def _load_recordings(self, filename):
105    recording_file = open(os.path.join(tempfile.gettempdir(), filename),
106                          'rb')
107    self._recordings = pickle.load(recording_file)
108    recording_file.close()
109
110  def _delete_recordings(self, filename):
111    full_path = os.path.join(tempfile.gettempdir(), filename)
112    if os.path.exists(full_path):
113      os.remove(full_path)
114
115  def _load_or_use_client(self, filename, http_client):
116    if os.path.exists(os.path.join(tempfile.gettempdir(), filename)):
117      self._load_recordings(filename)
118    else:
119      self.real_client = http_client
120
121  def use_cached_session(self, name=None, real_http_client=None):
122    """Attempts to load recordings from a previous live request.
123
124    If a temp file with the recordings exists, then it is used to fulfill
125    requests. If the file does not exist, then a real client is used to
126    actually make the desired HTTP requests. Requests and responses are
127    recorded and will be written to the desired temprary cache file when
128    close_session is called.
129
130    Args:
131      name: str (optional) The file name of session file to be used. The file
132            is loaded from the temporary directory of this machine. If no name
133            is passed in, a default name will be constructed using the
134            cache_name_prefix, cache_case_name, and cache_test_name of this
135            object.
136      real_http_client: atom.http_core.HttpClient the real client to be used
137                        if the cached recordings are not found. If the default
138                        value is used, this will be an
139                        atom.http_core.HttpClient.
140    """
141    if real_http_client is None:
142      real_http_client = atom.http_core.HttpClient()
143    if name is None:
144      self._recordings_cache_name = self.get_cache_file_name()
145    else:
146      self._recordings_cache_name = name
147    self._load_or_use_client(self._recordings_cache_name, real_http_client)
148
149  def close_session(self):
150    """Saves recordings in the temporary file named in use_cached_session."""
151    if self.real_client is not None:
152      self._save_recordings(self._recordings_cache_name)
153
154  def delete_session(self, name=None):
155    """Removes recordings from a previous live request."""
156    if name is None:
157      self._delete_recordings(self._recordings_cache_name)
158    else:
159      self._delete_recordings(name)
160
161  def get_cache_file_name(self):
162    return '%s.%s.%s' % (self.cache_name_prefix, self.cache_case_name,
163                         self.cache_test_name)
164
165  def _dump(self):
166    """Provides debug information in a string."""
167    output = 'MockHttpClient\n  real_client: %s\n  cache file name: %s\n' % (
168        self.real_client, self.get_cache_file_name())
169    output += '  recordings:\n'
170    i = 0
171    for recording in self._recordings:
172      output += '    recording %i is for: %s %s\n' % (
173          i, recording[0].method, str(recording[0].uri))
174      i += 1
175    return output
176
177
178def _match_request(http_request, stored_request):
179  """Determines whether a request is similar enough to a stored request
180     to cause the stored response to be returned."""
181  # Check to see if the host names match.
182  if (http_request.uri.host is not None
183      and http_request.uri.host != stored_request.uri.host):
184    return False
185  # Check the request path in the URL (/feeds/private/full/x)
186  elif http_request.uri.path != stored_request.uri.path:
187    return False
188  # Check the method used in the request (GET, POST, etc.)
189  elif http_request.method != stored_request.method:
190    return False
191  # If there is a gsession ID in either request, make sure that it is matched
192  # exactly.
193  elif ('gsessionid' in http_request.uri.query
194        or 'gsessionid' in stored_request.uri.query):
195    if 'gsessionid' not in stored_request.uri.query:
196      return False
197    elif 'gsessionid' not in http_request.uri.query:
198      return False
199    elif (http_request.uri.query['gsessionid']
200          != stored_request.uri.query['gsessionid']):
201      return False
202  # Ignores differences in the query params (?start-index=5&max-results=20),
203  # the body of the request, the port number, HTTP headers, just to name a
204  # few.
205  return True
206
207
208def _scrub_request(http_request):
209  """ Removes email address and password from a client login request.
210
211  Since the mock server saves the request and response in plantext, sensitive
212  information like the password should be removed before saving the
213  recordings. At the moment only requests sent to a ClientLogin url are
214  scrubbed.
215  """
216  if (http_request and http_request.uri and http_request.uri.path and
217      http_request.uri.path.endswith('ClientLogin')):
218    # Remove the email and password from a ClientLogin request.
219    http_request._body_parts = []
220    http_request.add_form_inputs(
221        {'form_data': 'client login request has been scrubbed'})
222  else:
223    # We can remove the body of the post from the recorded request, since
224    # the request body is not used when finding a matching recording.
225    http_request._body_parts = []
226  return http_request
227
228
229def _scrub_response(http_response):
230  return http_response
231
232
233class EchoHttpClient(object):
234  """Sends the request data back in the response.
235
236  Used to check the formatting of the request as it was sent. Always responds
237  with a 200 OK, and some information from the HTTP request is returned in
238  special Echo-X headers in the response. The following headers are added
239  in the response:
240  'Echo-Host': The host name and port number to which the HTTP connection is
241               made. If no port was passed in, the header will contain
242               host:None.
243  'Echo-Uri': The path portion of the URL being requested. /example?x=1&y=2
244  'Echo-Scheme': The beginning of the URL, usually 'http' or 'https'
245  'Echo-Method': The HTTP method being used, 'GET', 'POST', 'PUT', etc.
246  """
247
248  def request(self, http_request):
249    return self._http_request(http_request.uri, http_request.method,
250                              http_request.headers, http_request._body_parts)
251
252  def _http_request(self, uri, method, headers=None, body_parts=None):
253    body = StringIO.StringIO()
254    response = atom.http_core.HttpResponse(status=200, reason='OK', body=body)
255    if headers is None:
256      response._headers = {}
257    else:
258      # Copy headers from the request to the response but convert values to
259      # strings. Server response headers always come in as strings, so an int
260      # should be converted to a corresponding string when echoing.
261      for header, value in headers.iteritems():
262        response._headers[header] = str(value)
263    response._headers['Echo-Host'] = '%s:%s' % (uri.host, str(uri.port))
264    response._headers['Echo-Uri'] = uri._get_relative_path()
265    response._headers['Echo-Scheme'] = uri.scheme
266    response._headers['Echo-Method'] = method
267    for part in body_parts:
268      if isinstance(part, str):
269        body.write(part)
270      elif hasattr(part, 'read'):
271        body.write(part.read())
272    body.seek(0)
273    return response
274
275
276class SettableHttpClient(object):
277  """An HTTP Client which responds with the data given in set_response."""
278
279  def __init__(self, status, reason, body, headers):
280    """Configures the response for the server.
281
282    See set_response for details on the arguments to the constructor.
283    """
284    self.set_response(status, reason, body, headers)
285    self.last_request = None
286
287  def set_response(self, status, reason, body, headers):
288    """Determines the response which will be sent for each request.
289
290    Args:
291      status: An int for the HTTP status code, example: 200, 404, etc.
292      reason: String for the HTTP reason, example: OK, NOT FOUND, etc.
293      body: The body of the HTTP response as a string or a file-like
294            object (something with a read method).
295      headers: dict of strings containing the HTTP headers in the response.
296    """
297    self.response = atom.http_core.HttpResponse(status=status, reason=reason,
298        body=body)
299    self.response._headers = headers.copy()
300
301  def request(self, http_request):
302    self.last_request = http_request
303    return self.response
304
305
306class MockHttpResponse(atom.http_core.HttpResponse):
307
308  def __init__(self, status=None, reason=None, headers=None, body=None):
309    self._headers = headers or {}
310    if status is not None:
311      self.status = status
312    if reason is not None:
313      self.reason = reason
314    if body is not None:
315      # Instead of using a file-like object for the body, store as a string
316      # so that reads can be repeated.
317      if hasattr(body, 'read'):
318        self._body = body.read()
319      else:
320        self._body = body
321
322  def read(self):
323    return self._body