expose timeout on Connection class
This commit is contained in:
@@ -12,11 +12,12 @@ tracer.propagate = False
|
||||
class Connection(object):
|
||||
transport_schema = 'http'
|
||||
|
||||
def __init__(self, host='localhost', port=9200, url_prefix='', **kwargs):
|
||||
def __init__(self, host='localhost', port=9200, url_prefix='', timeout=10, **kwargs):
|
||||
self.host = '%s://%s:%s' % (self.transport_schema, host, port)
|
||||
if url_prefix:
|
||||
url_prefix = '/' + url_prefix.strip('/')
|
||||
self.url_prefix = url_prefix
|
||||
self.timeout = 10
|
||||
|
||||
def __repr__(self):
|
||||
return '<%s: %s>' % (self.__class__.__name__, self.host)
|
||||
@@ -56,14 +57,14 @@ class RequestsHttpConnection(Connection):
|
||||
super(RequestsHttpConnection, self).__init__(**kwargs)
|
||||
self.session = requests.session()
|
||||
|
||||
def perform_request(self, method, url, params=None, body=None):
|
||||
def perform_request(self, method, url, params=None, body=None, timeout=None):
|
||||
url = self.host + self.url_prefix + url
|
||||
|
||||
# use prepared requests so that requests formats url and params for us to log
|
||||
request = requests.Request(method, url, params=params or {}, data=body).prepare()
|
||||
start = time.time()
|
||||
try:
|
||||
response = self.session.send(request)
|
||||
response = self.session.send(request, timeout=timeout or self.timeout)
|
||||
duration = time.time() - start
|
||||
raw_data = response.text
|
||||
except (requests.ConnectionError, requests.Timeout) as e:
|
||||
|
||||
@@ -25,8 +25,9 @@ class TestRequestsConnection(TestCase):
|
||||
|
||||
self.assertEquals(1, connection.session.send.call_count)
|
||||
|
||||
timeout = kwargs.pop('timeout', connection.timeout)
|
||||
args, kwargs = connection.session.send.call_args
|
||||
self.assertEquals({}, kwargs)
|
||||
self.assertEquals({'timeout': timeout}, kwargs)
|
||||
self.assertEquals(1, len(args))
|
||||
return args[0]
|
||||
|
||||
@@ -115,7 +116,7 @@ class TestRequestsConnection(TestCase):
|
||||
@patch('elasticsearch.connection.tracer')
|
||||
def test_url_prefix(self, tracer):
|
||||
con = self._get_mock_connection({"url_prefix": "/some-prefix/"})
|
||||
request = self._get_request(con, 'GET', '/_search', body='{"answer": 42}')
|
||||
request = self._get_request(con, 'GET', '/_search', body='{"answer": 42}', timeout=0.1)
|
||||
|
||||
self.assertEquals('http://localhost:9200/some-prefix/_search', request.url)
|
||||
self.assertEquals('GET', request.method)
|
||||
|
||||
Reference in New Issue
Block a user