Honza Král
131 lines
3.4 KiB
Python
131 lines
3.4 KiB
Python
import time
|
|
import subprocess
|
|
import tempfile
|
|
import os
|
|
|
|
import requests
|
|
|
|
from elasticsearch import Elasticsearch
|
|
from elasticsearch.exceptions import ConnectionError, NotFoundError
|
|
|
|
from ..test_cases import TestCase, SkipTest
|
|
|
|
data_dir = None
|
|
|
|
CMD = """
|
|
elasticsearch \
|
|
-f \
|
|
-D es.cluster.name=%(cluster_name)s \
|
|
-D es.node.name=test_name \
|
|
-D es.http.port=%(port)s \
|
|
-D es.gateway.type=none \
|
|
-D es.index.store.type=memory \
|
|
-D es.discovery.zen.ping.multicast.enabled=false \
|
|
-D es.path.data=%(data_dir)s \
|
|
-D es.pidfile=%(pidfile)s \
|
|
>/dev/null 2>&1
|
|
"""
|
|
|
|
server = None
|
|
pidfile = tempfile.mktemp()
|
|
|
|
def get_client(**kwargs):
|
|
# 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'])
|
|
# update them with params
|
|
kw.update(kwargs)
|
|
|
|
# try and locate manual override in the local environment
|
|
try:
|
|
from test_elasticsearch.local import get_client as local_get_client
|
|
return local_get_client([os.environ['TEST_ES_SERVER']], **kw)
|
|
except ImportError:
|
|
# fallback to using vanilla client
|
|
return Elasticsearch([os.environ['TEST_ES_SERVER']], **kw)
|
|
|
|
|
|
def setup():
|
|
global server
|
|
|
|
# if use running ES instance, don't attempt to start our own
|
|
if 'TEST_ES_SERVER' not in os.environ:
|
|
# check installed
|
|
if subprocess.call('which elasticsearch >/dev/null 2>&1', shell=True) != 0:
|
|
raise SkipTest("No Elasticsearch server, skipping integration tests.")
|
|
|
|
args = {
|
|
'cluster_name': 'es_client_test',
|
|
'port': 9900,
|
|
'data_dir': tempfile.tempdir,
|
|
'pidfile': pidfile
|
|
}
|
|
|
|
# check running
|
|
try:
|
|
requests.get('http://localhost:%(port)s' % args)
|
|
except requests.ConnectionError:
|
|
pass
|
|
else:
|
|
raise SkipTest('Elasticsearch already running!')
|
|
|
|
|
|
cmd = CMD % args
|
|
|
|
server = subprocess.Popen(cmd, shell=True)
|
|
os.environ['TEST_ES_SERVER'] = 'localhost:%(port)s' % args
|
|
|
|
client = get_client()
|
|
|
|
# wait for yellow status
|
|
for _ in range(100):
|
|
time.sleep(.1)
|
|
try:
|
|
client.cluster.health(wait_for_status='yellow')
|
|
break
|
|
except ConnectionError:
|
|
continue
|
|
|
|
else:
|
|
# timeout
|
|
raise SkipTest("Elasticsearch failed to start.")
|
|
|
|
|
|
def teardown():
|
|
if server is not None:
|
|
with open(pidfile) as pidf:
|
|
pid = pidf.read()
|
|
os.kill(int(pid), 15)
|
|
server.wait()
|
|
|
|
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):
|
|
client = None
|
|
def setUp(self):
|
|
if ElasticTestCase.client is None:
|
|
ElasticTestCase.client = get_client()
|
|
self.client = ElasticTestCase.client
|
|
|
|
def tearDown(self):
|
|
self.client.indices.delete()
|
|
try:
|
|
self.client.indices.delete_template('*')
|
|
except NotFoundError:
|
|
pass
|
|
|
|
@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
|
|
|