Commits

Gregory Petukhov committed f729a51

Issue #90: make transport option more flexible

Comments (0)

Files changed (3)

     Public methods
     """
 
-    def __init__(self, response_body=None, transport='curl.CurlTransport',
+    def __init__(self, response_body=None, transport='grab.transport.curl.CurlTransport',
                  **kwargs):
         """
         Create Grab instance
         self.trigger_extensions('init')
         self._request_prepared = False
 
-        self.transport_name = transport
-        mod_name, cls_name = transport.split('.')
-        mod = __import__('grab.transport.%s' % mod_name, globals(),
-                         locals(), ['foo'])
-        self.transport = getattr(mod, cls_name)()
+        self.setup_transport(transport)
 
         self.reset()
         self.proxylist = None
         if response_body is not None:
             self.fake_response(response_body)
 
+    def setup_transport(self, transport_param):
+        self.transport_param = transport_param
+        if isinstance(transport_param, basestring):
+            mod_path, cls_name = transport_param.rsplit('.', 1)
+            mod = __import__(mod_path, globals(), locals(), ['foo'])
+            self.transport = getattr(mod, cls_name)()
+        elif callable(transport_param):
+            self.transport = transport_param()
+        else:
+            raise GrabMisuseError('Option `transport` should be string or callable. '\
+                                  'Got %s' % type(transport_param))
+
     def reset(self):
         """
         Reset all attributes which could be modified during previous request
         :param **kwargs: overrides settings of cloned grab instance
         """
 
-        g = Grab(transport=self.transport_name)
+        g = Grab(transport=self.transport_param)
         g.config = self.dump_config()
 
         if self.response is not None:
     'test.response_class',
     'test.charset_issue',
     'test.grab_pickle',
+    'test.grab_transport',
     # *** Tornado Test Server
     'test.tornado_server',
     # *** grab.tools
     parser = OptionParser()
     parser.add_option('-t', '--test', help='Run only specified tests')
     parser.add_option('--transport', help='Test specified transport',
-                      default='curl.CurlTransport')
+                      default='grab.transport.curl.CurlTransport')
     opts, args = parser.parse_args()
 
     test.util.GRAB_TRANSPORT = opts.transport

test/grab_transport.py

+from unittest import TestCase
+
+from grab import Grab
+from .tornado_util import SERVER
+from grab.transport.curl import CurlTransport
+import pickle
+
+class FakeTransport(CurlTransport):
+    def prepare_response(self, grab):
+        resp = super(FakeTransport, self).prepare_response(grab)
+        resp.body = 'Faked ' + resp.body
+        return resp
+
+
+def get_curl_transport():
+    return CurlTransport()
+
+
+def get_fake_transport():
+    return FakeTransport()
+
+
+class TestTransportTestCase(TestCase):
+    def setUp(self):
+        SERVER.reset()
+        SERVER.RESPONSE['get'] = 'XYZ'
+
+    def transport_option_logic(self, curl_transport, fake_transport):
+        g = Grab(transport=curl_transport)
+        g.go(SERVER.BASE_URL)
+        self.assertEqual(g.response.body, 'XYZ')
+
+        g2 = g.clone()
+        g.go(SERVER.BASE_URL)
+        self.assertEqual(g.response.body, 'XYZ')
+
+        g2_data = pickle.dumps(g2)
+        g3 = pickle.loads(g2_data)
+        g3.go(SERVER.BASE_URL)
+        self.assertEqual(g3.response.body, 'XYZ')
+
+        g = Grab(transport=fake_transport)
+        g.go(SERVER.BASE_URL)
+        self.assertEqual(g.response.body, 'Faked XYZ')
+
+        g2 = g.clone()
+        g.go(SERVER.BASE_URL)
+        self.assertEqual(g.response.body, 'Faked XYZ')
+
+        g2_data = pickle.dumps(g2)
+        g3 = pickle.loads(g2_data)
+        g3.go(SERVER.BASE_URL)
+        self.assertEqual(g3.response.body, 'Faked XYZ')
+
+    def test_transport_option_as_string(self):
+        self.transport_option_logic(
+            'grab.transport.curl.CurlTransport',
+            'test.grab_transport.FakeTransport',
+        )
+
+    def test_transport_option_as_class(self):
+        self.transport_option_logic(
+            CurlTransport,
+            FakeTransport,
+        )
+
+    def test_transport_option_as_function(self):
+        self.transport_option_logic(
+            get_curl_transport,
+            get_fake_transport,
+        )