diff --git a/elasticsearch/helpers/test.py b/elasticsearch/helpers/test.py new file mode 100644 index 00000000..d81f05aa --- /dev/null +++ b/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 + diff --git a/test_elasticsearch/README.rst b/test_elasticsearch/README.rst index 13f49b57..8e724879 100644 --- a/test_elasticsearch/README.rst +++ b/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`. diff --git a/test_elasticsearch/test_server/__init__.py b/test_elasticsearch/test_server/__init__.py index 68990242..8783c495 100644 --- a/test_elasticsearch/test_server/__init__.py +++ b/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() diff --git a/test_elasticsearch/test_server/test_common.py b/test_elasticsearch/test_server/test_common.py index bf61813f..dbe7c8d7 100644 --- a/test_elasticsearch/test_server/test_common.py +++ b/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'): diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index b3e432ec..e686de16 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/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): diff --git a/test_elasticsearch/test_server/test_memcached.py b/test_elasticsearch/test_server/test_memcached.py index 24e7061c..e9cf144d 100644 --- a/test_elasticsearch/test_server/test_memcached.py +++ b/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