Allow the same 'hosts' settings that Elasticsearch does

vrcmarcos · May 5, 2017 · 59eb6ea · 59eb6ea
1 parent 6761d82
commit 59eb6ea
Showing 1 changed file with 6 additions and 1 deletion.
diff --git 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
@@ -10,7 +11,11 @@
 
 
 def _get_elasticmock(hosts=None, *args, **kwargs):
-    elastic_key = 'localhost:9200' if hosts is None else '{0}:{1}'.format(hosts[0].get('host'), hosts[0].get('port'))
+    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: