|
35 | 35 | import json |
36 | 36 | import os |
37 | 37 | import urllib.parse |
| 38 | +from unittest import mock |
38 | 39 |
|
39 | 40 | from elasticsearch import VERSION as ES_VERSION |
40 | 41 | from elasticsearch import Elasticsearch |
@@ -126,6 +127,44 @@ def test_ping(instrument, elasticapm_client, elasticsearch): |
126 | 127 | assert span["context"]["http"]["status_code"] == 200 |
127 | 128 |
|
128 | 129 |
|
| 130 | +@pytest.mark.integrationtest |
| 131 | +def test_ping_db_instance(instrument, elasticapm_client, elasticsearch): |
| 132 | + # some ugly code to smuggle the "x-found-handling-cluster" header into |
| 133 | + # the response from Elasticsearch |
| 134 | + try: |
| 135 | + import elastic_transport |
| 136 | + |
| 137 | + pool = elasticsearch._transport.node_pool.get().pool |
| 138 | + except ImportError: |
| 139 | + pool = elasticsearch.transport.connection_pool.get_connection().pool |
| 140 | + |
| 141 | + original_urlopen = pool.urlopen |
| 142 | + |
| 143 | + def wrapper(*args, **kwargs): |
| 144 | + result = original_urlopen(*args, **kwargs) |
| 145 | + result.headers.update({"x-found-handling-cluster": "foo"}) |
| 146 | + return result |
| 147 | + |
| 148 | + elasticapm_client.begin_transaction("test") |
| 149 | + with mock.patch.object(pool, "urlopen", wraps=wrapper): |
| 150 | + result = elasticsearch.ping() |
| 151 | + |
| 152 | + elasticapm_client.end_transaction("test", "OK") |
| 153 | + parsed_url = urllib.parse.urlparse(os.environ["ES_URL"]) |
| 154 | + |
| 155 | + transaction = elasticapm_client.events[TRANSACTION][0] |
| 156 | + spans = elasticapm_client.spans_for_transaction(transaction) |
| 157 | + assert len(spans) == 1 |
| 158 | + span = spans[0] |
| 159 | + assert span["name"] == "ES HEAD /" |
| 160 | + assert span["context"]["destination"] == { |
| 161 | + "address": parsed_url.hostname, |
| 162 | + "port": parsed_url.port, |
| 163 | + "service": {"name": "", "resource": "elasticsearch/foo", "type": ""}, |
| 164 | + } |
| 165 | + assert span["context"]["db"]["instance"] == "foo" |
| 166 | + |
| 167 | + |
129 | 168 | @pytest.mark.integrationtest |
130 | 169 | def test_info(instrument, elasticapm_client, elasticsearch): |
131 | 170 | elasticapm_client.begin_transaction("test") |
|
0 commit comments