Provide helpers for running tests against elasticsearch

elasticsearch/helpers/test.py

@@ -0,0 +1,58 @@
+import time
+import os
+try:
+    # python 2.6
+    from unittest2 import TestCase, SkipTest
+except ImportError:
+    from unittest import TestCase, SkipTest
+
+from elasticsearch import Elasticsearch
+from elasticsearch.exceptions import ConnectionError
+
+def get_test_client():
+    # construct kwargs from the environment
+    kw = {}
+    if 'TEST_ES_CONNECTION' in os.environ:
+        from elasticsearch import connection
+        kw['connection_class'] = getattr(connection, os.environ['TEST_ES_CONNECTION'])
+
+    client = Elasticsearch([os.environ.get('TEST_ES_SERVER', {})], **kw)
+
+    # wait for yellow status
+    for _ in range(100):
+        time.sleep(.1)
+        try:
+            client.cluster.health(wait_for_status='yellow')
+            return client
+        except ConnectionError:
+            continue
+    else:
+        # timeout
+        raise SkipTest("Elasticsearch failed to start.")
+
+def _get_version(version_string):
+    version = version_string.strip().split('.')
+    return tuple(int(v) if v.isdigit() else 999 for v in version)
+
+class ElasticsearchTestCase(TestCase):
+    @staticmethod
+    def _get_client():
+        return get_test_client()
+
+    @classmethod
+    def setUpClass(cls):
+        super(ElasticsearchTestCase, cls).setUpClass()
+        cls.client = cls._get_client()
+
+    def tearDown(self):
+        super(ElasticsearchTestCase, self).tearDown()
+        self.client.indices.delete(index='*')
+        self.client.indices.delete_template(name='*', ignore=404)
+
+    @property
+    def es_version(self):
+        if not hasattr(self, '_es_version'):
+            version_string = self.client.info()['version']['number']
+            self._es_version = _get_version(version_string)
+        return self._es_version
+

test_elasticsearch/README.rst

@@ -43,12 +43,10 @@ Customizing the tests
 ---------------------
 
 You can create a `local.py` file in the `test_elasticsearch` directory which
-should contain a `get_client` function:
+should contain a `get_client` function.
 
-    def get_client(hosts, ** kwargs):
-        ...
-
-If this file exists the function will be used instead of the built in one to
-construct the client used for any integration tests. You can use this to make
-sure your plugins and extensions work with `elasticsearch-py`.
+If this file exists the function will be used instead of
+`elasticsearch.helpers.test.get_test_client` to construct the client used for
+any integration tests. You can use this to make sure your plugins and
+extensions work with `elasticsearch-py`.
 

test_elasticsearch/test_server/__init__.py

@@ -1,10 +1,4 @@
-import time
-import os
-
-from elasticsearch import Elasticsearch
-from elasticsearch.exceptions import ConnectionError, NotFoundError
-
-from ..test_cases import TestCase, SkipTest
+from elasticsearch.helpers.test import get_test_client, ElasticsearchTestCase as BaseTestCase
 
 client = None
 
@@ -13,54 +7,21 @@ def get_client():
     if client is not None:
         return client
 
-    # construct kwargs from the environment
-    kw = {}
-    if 'TEST_ES_CONNECTION' in os.environ:
-        from elasticsearch import connection
-        kw['connection_class'] = getattr(connection, os.environ['TEST_ES_CONNECTION'])
-
     # try and locate manual override in the local environment
     try:
         from test_elasticsearch.local import get_client as local_get_client
-        client = local_get_client([os.environ.get('TEST_ES_SERVER', {})], **kw)
+        client = local_get_client()
     except ImportError:
         # fallback to using vanilla client
-        client = Elasticsearch([os.environ.get('TEST_ES_SERVER', {})], **kw)
+        client = get_test_client()
+
+    return client
 
-    # wait for yellow status
-    for _ in range(100):
-        time.sleep(.1)
-        try:
-            client.cluster.health(wait_for_status='yellow')
-            return client
-        except ConnectionError:
-            continue
-    else:
-        # timeout
-        raise SkipTest("Elasticsearch failed to start.")
 
 def setup():
     get_client()
 
-ES_VERSION = None
-
-def _get_version(version_string):
-    version = version_string.strip().split('.')
-    return tuple(int(v) if v.isdigit() else 999 for v in version)
-
-class ElasticTestCase(TestCase):
-    def setUp(self):
-        self.client = get_client()
-
-    def tearDown(self):
-        self.client.indices.delete(index='*')
-        self.client.indices.delete_template(name='*', ignore=404)
-
-    @property
-    def es_version(self):
-        global ES_VERSION
-        if ES_VERSION is None:
-            version_string = self.client.info()['version']['number']
-            ES_VERSION = _get_version(version_string)
-        return ES_VERSION
-
+class ElasticsearchTestCase(BaseTestCase):
+    @staticmethod
+    def _get_client():
+        return get_client()

test_elasticsearch/test_server/test_common.py

@@ -10,9 +10,10 @@ import yaml
 
 from elasticsearch import TransportError
 from elasticsearch.compat import string_types
+from elasticsearch.helpers.test import _get_version
 
 from ..test_cases import SkipTest
-from . import ElasticTestCase, _get_version
+from . import ElasticsearchTestCase
 
 # some params had to be changed in python, keep track of them so we can rename
 # those in the tests accordingly
@@ -33,7 +34,7 @@ IMPLEMENTED_FEATURES = ('regex', 'gtelte')
 class InvalidActionType(Exception):
     pass
 
-class YamlTestCase(ElasticTestCase):
+class YamlTestCase(ElasticsearchTestCase):
     def setUp(self):
         super(YamlTestCase, self).setUp()
         if hasattr(self, '_setup_code'):

test_elasticsearch/test_server/test_helpers.py

@@ -1,9 +1,9 @@
 from elasticsearch import helpers
 
-from . import ElasticTestCase
+from . import ElasticsearchTestCase
 from ..test_cases import SkipTest
 
-class TestStreamingBulk(ElasticTestCase):
+class TestStreamingBulk(ElasticsearchTestCase):
     def test_actions_remain_unchanged(self):
         actions = [{'_id': 1}, {'_id': 2}]
         for ok, item in helpers.streaming_bulk(self.client, actions, index='test-index', doc_type='answers'):
@@ -53,7 +53,7 @@ class TestStreamingBulk(ElasticTestCase):
         self.assertEquals({'f': 'v'}, self.client.get(index='i', id=47)['_source'])
 
 
-class TestBulk(ElasticTestCase):
+class TestBulk(ElasticsearchTestCase):
     def test_all_documents_get_inserted(self):
         docs = [{"answer": x, '_id': x} for x in range(100)]
         success, failed = helpers.bulk(self.client, docs, index='test-index', doc_type='answers', refresh=True)
@@ -128,7 +128,7 @@ class TestBulk(ElasticTestCase):
         self.assertEquals(1, failed)
 
 
-class TestScan(ElasticTestCase):
+class TestScan(ElasticsearchTestCase):
     def test_all_documents_are_read(self):
         bulk = []
         for x in range(100):
@@ -142,7 +142,7 @@ class TestScan(ElasticTestCase):
         self.assertEquals(set(map(str, range(100))), set(d['_id'] for d in docs))
         self.assertEquals(set(range(100)), set(d['_source']['answer'] for d in docs))
 
-class TestReindex(ElasticTestCase):
+class TestReindex(ElasticsearchTestCase):
     def test_all_documents_get_moved(self):
         bulk = []
         for x in range(100):

test_elasticsearch/test_server/test_memcached.py

@@ -4,10 +4,10 @@ from elasticsearch import Elasticsearch, MemcachedConnection, NotFoundError
 from elasticsearch.transport import ADDRESS_RE
 from elasticsearch.compat import u
 
-from . import ElasticTestCase
+from . import ElasticsearchTestCase
 from ..test_cases import SkipTest
 
-class TestMemcachedConnection(ElasticTestCase):
+class TestMemcachedConnection(ElasticsearchTestCase):
     def setUp(self):
         try:
             import pylibmc