diff --git a/elasticsearch/client.py b/elasticsearch/client.py index 0dd85e54..1d9b9cd6 100644 --- a/elasticsearch/client.py +++ b/elasticsearch/client.py @@ -83,7 +83,7 @@ class Elasticsearch(object): Elasticsearch low-level client. Provides a straightforward mapping from Python to ES REST endpoints. """ - def __init__(self, hosts=None, **kwargs): + def __init__(self, hosts=None, transport_class=Transport, **kwargs): """ :arg hosts: list of nodes we should connect to. Node should be a dictionary ({"host": "localhost", "port": 9200}), the entire dictionary @@ -92,11 +92,13 @@ class Elasticsearch(object): translated to a dictionary automatically. If no value is given the :class:`~elasticsearch.Connection` class defaults will be used. + :arg transport_class: :class:`~elasticsearch.Transport` subclass to use. + :arg kwargs: any additional arguments will be passed on to the :class:`~elasticsearch.Transport` class and, subsequently, to the :class:`~elasticsearch.Connection` instances. """ - self.transport = Transport(_normalize_hosts(hosts), **kwargs) + self.transport = transport_class(_normalize_hosts(hosts), **kwargs) # namespaced clients for compatibility with API names self.indices = InidicesClient(self) diff --git a/test_elasticsearch/test_cases.py b/test_elasticsearch/test_cases.py new file mode 100644 index 00000000..ffeb5a81 --- /dev/null +++ b/test_elasticsearch/test_cases.py @@ -0,0 +1,47 @@ +from unittest import TestCase +from collections import defaultdict + +from elasticsearch import Elasticsearch + +class DummyTransport(object): + def __init__(self, hosts, responses=None, **kwargs): + self.responses = responses + self.call_count = 0 + self.calls = defaultdict(list) + + def perform_request(self, method, url, params=None, body=None): + resp = 200, {} + if self.responses: + resp = self.responses[self.call_count] + self.call_count += 1 + self.calls[(method, url)].append((params, body)) + return resp + + +class ElasticsearchTestCase(TestCase): + def setUp(self): + super(ElasticsearchTestCase, self).setUp() + self.client = Elasticsearch(transport_class=DummyTransport) + + def assert_call_count_equals(self, count): + self.assertEquals(count, self.client.transport.call_count) + + def assert_url_called(self, method, url, count=1): + self.assertIn((method, url), self.client.transport.calls) + calls = self.client.transport.calls[(method, url)] + self.assertEquals(count, len(calls)) + return calls + + +class TestElasticsearchTestCase(ElasticsearchTestCase): + def test_our_transport_used(self): + self.assertIsInstance(self.client.transport, DummyTransport) + + def test_start_with_0_call(self): + self.assert_call_count_equals(0) + + def test_each_call_is_recorded(self): + self.client.transport.perform_request('GET', '/') + self.client.transport.perform_request('DELETE', '/42', params={}, body='body') + self.assert_call_count_equals(2) + self.assertEquals([({}, 'body')], self.assert_url_called('DELETE', '/42', 1))