From 2576878f2eba16d8ff1594849283bfe26a7e791a Mon Sep 17 00:00:00 2001 From: Kevin McCarthy Date: Sat, 30 Jun 2012 15:10:20 -1000 Subject: [PATCH] Cassettes now store multiple requests --- README.md | 2 +- test.py | 17 +++++++++++++++-- vcr/cassette.py | 23 +++++++++++++++++------ vcr/files.py | 7 +++++-- vcr/patch.py | 3 ++- vcr/stubs.py | 24 +++++++++++++++++++----- 6 files changed, 59 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 5a0cc29..8d8e132 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,8 @@ Ruby's VCR are not compatible with VCR.py. The API is similar but VCR.py doesn't have nearly as many features. ##Known Issues - * Only works with the first HTTP request made in each use_cassette block. That's because I haven't gotten around to implementing multiple requests yet. This means anything that does a 301 / 302 won't work. * Probably only works with urllib2. Hey, doesn't [requests](http://docs.python-requests.org/en/latest/index.html) use urllib2? Maybe that works too then. + * Loading urls that do a 301/302 redirect don't work ##Similar libraries in Python Neither of these really implement the API I want, but I have cribbed some code diff --git a/test.py b/test.py index 1379fa9..ea3d5a7 100644 --- a/test.py +++ b/test.py @@ -1,6 +1,7 @@ import os import unittest import vcr +from vcr.cassette import Cassette import urllib2 TEST_CASSETTE_FILE = 'test/test_req.yaml' @@ -48,7 +49,7 @@ class TestHttps(unittest.TestCase): os.remove(TEST_CASSETTE_FILE) except OSError: pass - + def test_response_code(self): code = urllib2.urlopen('https://api.twitter.com/1/legal/tos.json').getcode() with vcr.use_cassette(TEST_CASSETTE_FILE): @@ -65,7 +66,19 @@ class TestHttps(unittest.TestCase): with vcr.use_cassette(TEST_CASSETTE_FILE): headers = urllib2.urlopen('https://api.twitter.com/1/legal/tos.json').info().items() self.assertEqual(headers, urllib2.urlopen('https://api.twitter.com/1/legal/tos.json').info().items()) - + +class TestCassette(unittest.TestCase): + def test_serialize_cassette(self): + c1 = Cassette() + c1.requests = ['a','b','c'] + c1.responses = ['d','e','f'] + ser = c1.serialize() + c2 = Cassette(ser) + self.assertEqual(c1.requests,c2.requests) + self.assertEqual(c1.responses,c2.responses) + + + if __name__ == '__main__': unittest.main() diff --git a/vcr/cassette.py b/vcr/cassette.py index d2f1a2e..e101fba 100644 --- a/vcr/cassette.py +++ b/vcr/cassette.py @@ -1,16 +1,27 @@ -import yaml - - class Cassette(object): - def __init__(self): + def __init__(self, ser_cassette=None): self.requests = [] self.responses = [] + if ser_cassette: + self._unserialize(ser_cassette) def serialize(self): - return yaml.dump([{ + return ([{ 'request': req, 'response': res, - } for req,res in zip(self.requests,self.responses)]) + } for req, res in zip(self.requests, self.responses)]) + def _unserialize(self, source): + self.requests, self.responses = [r['request'] for r in source], [r['response'] for r in source] + def get_request(self, match): + try: + return self.requests[self.requests.index(match)] + except ValueError: + return None + def get_response(self, match): + try: + return self.responses[self.requests.index(match)] + except ValueError: + return None diff --git a/vcr/files.py b/vcr/files.py index 89220ec..79047d7 100644 --- a/vcr/files.py +++ b/vcr/files.py @@ -1,10 +1,13 @@ import os import yaml +from .cassette import Cassette def load_cassette(cassette_path): try: - return yaml.load(open(cassette_path)) + pc = yaml.load(open(cassette_path)) + cassette = Cassette(pc) + return cassette except IOError: return None @@ -14,4 +17,4 @@ def save_cassette(cassette_path, cassette): if not os.path.exists(dirname): os.makedirs(dirname) with open(cassette_path, 'wc') as cassette_file: - cassette_file.write(cassette.serialize()) + cassette_file.write(yaml.dump(cassette.serialize())) diff --git a/vcr/patch.py b/vcr/patch.py index 4f26307..ce47b3a 100644 --- a/vcr/patch.py +++ b/vcr/patch.py @@ -12,14 +12,15 @@ def install(cassette_path): httplib.HTTPConnection._vcr_cassette_path = cassette_path httplib.HTTPSConnection._vcr_cassette_path = cassette_path + def reset(): httplib.HTTPConnection = httplib.HTTP._connection_class = _HTTPConnection httplib.HTTPSConnection = httplib.HTTPS._connection_class = \ _HTTPSConnection + @contextmanager def use_cassette(cassette_path): install(cassette_path) yield reset() - diff --git a/vcr/stubs.py b/vcr/stubs.py index 5f87cca..a01b33d 100644 --- a/vcr/stubs.py +++ b/vcr/stubs.py @@ -28,9 +28,23 @@ class VCRHTTPConnection(HTTPConnection): def _save_cassette(self): save_cassette(self._vcr_cassette_path, self._cassette) - def request(self, method, url, body=None, headers={}): + def _load_old_response(self): old_cassette = load_cassette(self._vcr_cassette_path) if old_cassette: + return old_cassette.get_response(self._vcr) + + def request(self, method, url, body=None, headers={}): + """ + Persist the request metadata in self._vcr + """ + self._vcr = { + 'method': method, + 'url': url, + 'body': body, + 'headers': headers, + } + old_cassette = load_cassette(self._vcr_cassette_path) + if old_cassette and old_cassette.get_request(self._vcr): return self._cassette.requests.append(dict( method=method, @@ -41,8 +55,8 @@ class VCRHTTPConnection(HTTPConnection): return HTTPConnection.request(self, method, url, body=body, headers=headers) def getresponse(self, buffering=False): - old_cassette = load_cassette(self._vcr_cassette_path) - if not old_cassette: + old_response = self._load_old_response() + if not old_response: response = HTTPConnection.getresponse(self) self._cassette.responses.append({ 'status': {'code': response.status, 'message': response.reason}, @@ -50,8 +64,8 @@ class VCRHTTPConnection(HTTPConnection): 'body': {'string': response.read()}, }) self._save_cassette() - old_cassette = load_cassette(self._vcr_cassette_path) - return VCRHTTPResponse(old_cassette[0]['response']) + old_response = self._load_old_response() + return VCRHTTPResponse(old_response) class VCRHTTPSConnection(VCRHTTPConnection):