diff --git a/elasticmock/__init__.py b/elasticmock/__init__.py index 0b92193..699587b 100644 --- a/elasticmock/__init__.py +++ b/elasticmock/__init__.py @@ -2,6 +2,7 @@ from functools import wraps +from elasticsearch.client import _normalize_hosts from mock import patch from elasticmock.fake_elasticsearch import FakeElasticsearch @@ -9,8 +10,12 @@ ELASTIC_INSTANCES = {} -def _get_elasticmock(hosts=None): - elastic_key = 'localhost:9200' if hosts is None else '{0}:{1}'.format(hosts[0].get('host'), hosts[0].get('port')) +def _get_elasticmock(hosts=None, *args, **kwargs): + host = _normalize_hosts(hosts)[0] + elastic_key = '{0}:{1}'.format( + host.get('host', 'localhost'), host.get('port', 9200) + ) + if elastic_key in ELASTIC_INSTANCES: connection = ELASTIC_INSTANCES.get(elastic_key) else: