Skip to content

Commit

Permalink
WIP elasticsearch: tests against elasticsearch 8
Browse files Browse the repository at this point in the history
  • Loading branch information
xrmx committed Apr 29, 2024
1 parent 3291f38 commit 4b3f945
Show file tree
Hide file tree
Showing 4 changed files with 94 additions and 38 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,8 @@ def _instrument(self, **kwargs):

def _uninstrument(self, **kwargs):
# pylint: disable=no-member
unwrap(elasticsearch.Transport, "perform_request")
transport_class = elastic_transport.Transport if es_transport_split else elasticsearch.Transport
unwrap(transport_class, "perform_request")


_regex_doc_url = re.compile(r"/_doc/([^/]+)")
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
asgiref==3.7.2
attrs==23.2.0
Deprecated==1.2.14
elasticsearch==8.12.1
elasticsearch-dsl==8.12.0
elastic-transport==8.12.0
importlib-metadata==7.1.0
iniconfig==2.0.0
packaging==23.2
pluggy==1.4.0
py==1.11.0
py-cpuinfo==9.0.0
pytest==7.1.3
pytest-benchmark==4.0.0
python-dateutil==2.8.2
six==1.16.0
tomli==2.0.1
typing_extensions==4.10.0
urllib3==2.2.1
wrapt==1.16.0
zipp==3.17.0
-e opentelemetry-instrumentation
-e instrumentation/opentelemetry-instrumentation-elasticsearch
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,11 @@

major_version = elasticsearch.VERSION[0]

NodeAPIResponse, ApiResponseMeta, HttpHeaders = None, None, None

if major_version == 8:
from elastic_transport._node import NodeApiResponse
from elastic_transport import ApiResponseMeta, HttpHeaders
from . import helpers_es8 as helpers # pylint: disable=no-name-in-module
elif major_version == 7:
from . import helpers_es7 as helpers # pylint: disable=no-name-in-module
Expand All @@ -51,25 +55,41 @@


def normalize_arguments(doc_type, body=None):
if major_version == 7:
return {"document": body} if body else {}
return (
{"body": body, "doc_type": doc_type}
if body
else {"doc_type": doc_type}
)
if major_version < 7:
return (
{"body": body, "doc_type": doc_type}
if body
else {"doc_type": doc_type}
)
return {"document": body} if body else {}


def get_elasticsearch_client(*args, **kwargs):
client = Elasticsearch(*args, **kwargs)
if major_version == 7:
if major_version == 8:
client._verified_elasticsearch = True
elif major_version == 7:
client.transport._verified_elasticsearch = True
return client


@mock.patch(
"elasticsearch.connection.http_urllib3.Urllib3HttpConnection.perform_request"
)
class PerformRequestMock:
@staticmethod
def path():
if major_version < 8:
return "elasticsearch.connection.http_urllib3.Urllib3HttpConnection.perform_request"
else:
return "elastic_transport._node._http_urllib3.Urllib3HttpNode.perform_request"

@staticmethod
def response(body: str):
if major_version < 8:
return (1, {}, body)
else:
return NodeApiResponse(ApiResponseMeta(status=200, headers=HttpHeaders({}), duration=100, http_version="1.1", node="node"), body.encode())


@mock.patch(PerformRequestMock.path())
class TestElasticsearchIntegration(TestBase):
search_attributes = {
SpanAttributes.DB_SYSTEM: "elasticsearch",
Expand All @@ -96,7 +116,7 @@ def tearDown(self):
ElasticsearchInstrumentor().uninstrument()

def test_instrumentor(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")

es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.index(
Expand Down Expand Up @@ -147,7 +167,7 @@ def test_prefix_arg(self, request_mock):
prefix = "prefix-from-env"
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor(span_name_prefix=prefix).instrument()
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
self._test_prefix(prefix)

def test_prefix_env(self, request_mock):
Expand All @@ -156,7 +176,7 @@ def test_prefix_env(self, request_mock):
os.environ[env_var] = prefix
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument()
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
del os.environ[env_var]
self._test_prefix(prefix)

Expand All @@ -174,10 +194,8 @@ def _test_prefix(self, prefix):
self.assertTrue(span.name.startswith(prefix))

def test_result_values(self, request_mock):
request_mock.return_value = (
1,
{},
'{"found": false, "timed_out": true, "took": 7}',
request_mock.return_value = PerformRequestMock.response(
'{"found": false, "timed_out": true, "took": 7}'
)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.get(
Expand All @@ -201,8 +219,11 @@ def test_trace_error_unknown(self, request_mock):

def test_trace_error_not_found(self, request_mock):
msg = "record not found"
exc = elasticsearch.exceptions.NotFoundError(404, msg)
request_mock.return_value = (1, {}, "{}")
if major_version == 8:
exc = elasticsearch.exceptions.NotFoundError(404, msg, body=None)
else:
exc = elasticsearch.exceptions.NotFoundError(404, msg)
request_mock.return_value = PerformRequestMock.response("{}")
request_mock.side_effect = exc
self._test_trace_error(StatusCode.ERROR, exc)

Expand All @@ -227,7 +248,7 @@ def _test_trace_error(self, code, exc):
)

def test_parent(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
with self.tracer.start_as_current_span("parent"):
es.index(
Expand All @@ -245,7 +266,7 @@ def test_parent(self, request_mock):
self.assertEqual(child.parent.span_id, parent.context.span_id)

def test_multithread(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
ev = threading.Event()

Expand Down Expand Up @@ -292,7 +313,7 @@ def target2():
self.assertIsNone(s3.parent)

def test_dsl_search(self, request_mock):
request_mock.return_value = (1, {}, '{"hits": {"hits": []}}')
request_mock.return_value = PerformRequestMock.response('{"hits": {"hits": []}}')

client = get_elasticsearch_client(hosts=["http://localhost:9200"])
search = Search(using=client, index="test-index").filter(
Expand All @@ -310,7 +331,7 @@ def test_dsl_search(self, request_mock):
)

def test_dsl_search_sanitized(self, request_mock):
request_mock.return_value = (1, {}, '{"hits": {"hits": []}}')
request_mock.return_value = PerformRequestMock.response('{"hits": {"hits": []}}')
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
search = Search(using=client, index="test-index").filter(
"term", author="testing"
Expand All @@ -327,7 +348,7 @@ def test_dsl_search_sanitized(self, request_mock):
)

def test_dsl_create(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
Article.init(using=client)

Expand All @@ -354,7 +375,7 @@ def test_dsl_create(self, request_mock):
)

def test_dsl_create_sanitized(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
Article.init(using=client)

Expand All @@ -370,7 +391,7 @@ def test_dsl_create_sanitized(self, request_mock):
)

def test_dsl_index(self, request_mock):
request_mock.return_value = (1, {}, helpers.dsl_index_result[2])
request_mock.return_value = PerformRequestMock.response(helpers.dsl_index_result[2])

client = get_elasticsearch_client(hosts=["http://localhost:9200"])
article = Article(
Expand Down Expand Up @@ -404,6 +425,9 @@ def test_request_hook(self, request_mock):
request_hook_kwargs_attribute = "request_hook.kwargs"

def request_hook(span, method, url, kwargs):
# FIXME: this is done only to get tests passing, need a more clueful solution
if major_version == 8:
kwargs = dict(kwargs["headers"])
attributes = {
request_hook_method_attribute: method,
request_hook_url_attribute: url,
Expand All @@ -416,10 +440,8 @@ def request_hook(span, method, url, kwargs):
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument(request_hook=request_hook)

request_mock.return_value = (
1,
{},
'{"found": false, "timed_out": true, "took": 7}',
request_mock.return_value = PerformRequestMock.response(
'{"found": false, "timed_out": true, "took": 7}'
)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
index = "test-index"
Expand All @@ -439,12 +461,17 @@ def request_hook(span, method, url, kwargs):
"GET", spans[0].attributes[request_hook_method_attribute]
)
expected_url = f"/{index}/_doc/{doc_id}"
if major_version == 8:
expected_url += "?realtime=true&refresh=true"
self.assertEqual(
expected_url,
spans[0].attributes[request_hook_url_attribute],
)

if major_version == 7:
if major_version == 8:
# FIXME: kwargs passed to request_hook on 8 are completely different
expected_kwargs = {"accept": "application/vnd.elasticsearch+json; compatible-with=8"}
elif major_version == 7:
expected_kwargs = {
**kwargs,
"headers": {"accept": "application/json"},
Expand All @@ -461,6 +488,9 @@ def test_response_hook(self, request_mock):

def response_hook(span, response):
if span and span.is_recording():
# FIXME: something more clean
if major_version == 8:
response = response.body
span.set_attribute(
response_attribute_name, json.dumps(response)
)
Expand Down Expand Up @@ -492,7 +522,7 @@ def response_hook(span, response):
},
}

request_mock.return_value = (1, {}, json.dumps(response_payload))
request_mock.return_value = PerformRequestMock.response(json.dumps(response_payload))
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.get(
index="test-index", **normalize_arguments(doc_type="_doc"), id=1
Expand All @@ -512,7 +542,7 @@ def test_no_op_tracer_provider(self, request_mock):
tracer_provider=trace.NoOpTracerProvider()
)
response_payload = '{"found": false, "timed_out": true, "took": 7}'
request_mock.return_value = (1, {}, response_payload)
request_mock.return_value = PerformRequestMock.response(response_payload)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
res = es.get(
index="test-index", **normalize_arguments(doc_type="_doc"), id=1
Expand Down Expand Up @@ -543,7 +573,7 @@ def test_body_sanitization(self, _):
)

def test_bulk(self, request_mock):
request_mock.return_value = (1, {}, "{}")
request_mock.return_value = PerformRequestMock.response("{}")

es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.bulk(
Expand Down
6 changes: 4 additions & 2 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,9 @@ envlist =
; below mean these dependencies are being used:
; 0: elasticsearch-dsl==6.4.0 elasticsearch==6.8.2
; 1: elasticsearch-dsl==7.4.1 elasticsearch==7.17.9
py3{8,9,10,11}-test-instrumentation-elasticsearch-{0,1}
pypy3-test-instrumentation-elasticsearch-{0,1}
; 2: elasticsearch-dsl>=8.0,<8.13 elasticsearch>=8.0,<8.13
py3{8,9,10,11}-test-instrumentation-elasticsearch-{0,1,2}
pypy3-test-instrumentation-elasticsearch-{0,1,2}

; opentelemetry-instrumentation-falcon
; py310 does not work with falcon 1
Expand Down Expand Up @@ -645,6 +646,7 @@ commands_pre =
elasticsearch: pip install opentelemetry-test-utils@{env:CORE_REPO}\#egg=opentelemetry-test-utils&subdirectory=tests/opentelemetry-test-utils
elasticsearch-0: pip install -r {toxinidir}/instrumentation/opentelemetry-instrumentation-elasticsearch/test-requirements-0.txt
elasticsearch-1: pip install -r {toxinidir}/instrumentation/opentelemetry-instrumentation-elasticsearch/test-requirements-1.txt
elasticsearch-2: pip install -r {toxinidir}/instrumentation/opentelemetry-instrumentation-elasticsearch/test-requirements-2.txt

asyncio: pip install opentelemetry-api@{env:CORE_REPO}\#egg=opentelemetry-api&subdirectory=opentelemetry-api
asyncio: pip install opentelemetry-semantic-conventions@{env:CORE_REPO}\#egg=opentelemetry-semantic-conventions&subdirectory=opentelemetry-semantic-conventions
Expand Down

0 comments on commit 4b3f945

Please sign in to comment.