Skip to content

Commit 8a6b408

Browse files
codebotentoumorokoshi
codeboten
authored and
toumorokoshi
committed
Using InMemorySpanExporter for wsgi/flask tests (#306)
The InMemorySpanExporter provides a friendly interface to retrieving span information, reducing the need for mocking in unit tests. Signed-off-by: Alex Boten <aboten@lightstep.com>
1 parent b01f7e8 commit 8a6b408

4 files changed

Lines changed: 90 additions & 120 deletions

File tree

‎ext/opentelemetry-ext-flask/tests/test_flask_integration.py‎

Lines changed: 54 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# limitations under the License.
1414

1515
import unittest
16-
from unittest import mock
1716

1817
from flask import Flask
1918
from werkzeug.test import Client
@@ -24,18 +23,28 @@
2423
from opentelemetry.ext.testutil.wsgitestutil import WsgiTestBase
2524

2625

26+
def expected_attributes(override_attributes):
27+
default_attributes = {
28+
"component": "http",
29+
"http.method": "GET",
30+
"http.server_name": "localhost",
31+
"http.scheme": "http",
32+
"host.port": 80,
33+
"http.host": "localhost",
34+
"http.target": "/",
35+
"http.flavor": "1.1",
36+
"http.status_text": "OK",
37+
"http.status_code": 200,
38+
}
39+
for key, val in override_attributes.items():
40+
default_attributes[key] = val
41+
return default_attributes
42+
43+
2744
class TestFlaskIntegration(WsgiTestBase):
2845
def setUp(self):
2946
super().setUp()
3047

31-
self.span_attrs = {}
32-
33-
def setspanattr(key, value):
34-
self.assertIsInstance(key, str)
35-
self.span_attrs[key] = value
36-
37-
self.span.set_attribute = setspanattr
38-
3948
self.app = Flask(__name__)
4049

4150
def hello_endpoint(helloid):
@@ -49,100 +58,57 @@ def hello_endpoint(helloid):
4958
self.client = Client(self.app, BaseResponse)
5059

5160
def test_simple(self):
52-
resp = self.client.get("/hello/123")
53-
self.assertEqual(200, resp.status_code)
54-
self.assertEqual([b"Hello: 123"], list(resp.response))
55-
56-
self.start_span.assert_called_with(
57-
"hello_endpoint",
58-
trace_api.INVALID_SPAN_CONTEXT,
59-
kind=trace_api.SpanKind.SERVER,
60-
attributes={
61-
"component": "http",
62-
"http.method": "GET",
63-
"http.server_name": "localhost",
64-
"http.scheme": "http",
65-
"host.port": 80,
66-
"http.host": "localhost",
61+
expected_attrs = expected_attributes(
62+
{
6763
"http.target": "/hello/123",
68-
"http.flavor": "1.1",
6964
"http.route": "/hello/<int:helloid>",
70-
},
71-
start_time=mock.ANY,
72-
)
73-
74-
# TODO: Change this test to use the SDK, as mocking becomes painful
75-
76-
self.assertEqual(
77-
self.span_attrs,
78-
{"http.status_code": 200, "http.status_text": "OK"},
65+
}
7966
)
67+
resp = self.client.get("/hello/123")
68+
self.assertEqual(200, resp.status_code)
69+
self.assertEqual([b"Hello: 123"], list(resp.response))
70+
span_list = self.memory_exporter.get_finished_spans()
71+
self.assertEqual(len(span_list), 1)
72+
self.assertEqual(span_list[0].name, "hello_endpoint")
73+
self.assertEqual(span_list[0].kind, trace_api.SpanKind.SERVER)
74+
self.assertEqual(span_list[0].attributes, expected_attrs)
8075

8176
def test_404(self):
82-
resp = self.client.post("/bye")
83-
self.assertEqual(404, resp.status_code)
84-
resp.close()
85-
86-
self.start_span.assert_called_with(
87-
"/bye",
88-
trace_api.INVALID_SPAN_CONTEXT,
89-
kind=trace_api.SpanKind.SERVER,
90-
attributes={
91-
"component": "http",
77+
expected_attrs = expected_attributes(
78+
{
9279
"http.method": "POST",
93-
"http.server_name": "localhost",
94-
"http.scheme": "http",
95-
"host.port": 80,
96-
"http.host": "localhost",
9780
"http.target": "/bye",
98-
"http.flavor": "1.1",
99-
},
100-
start_time=mock.ANY,
81+
"http.status_text": "NOT FOUND",
82+
"http.status_code": 404,
83+
}
10184
)
10285

103-
# Nope, this uses Tracer.use_span(end_on_exit)
104-
# self.assertEqual(1, self.span.end.call_count)
105-
# TODO: Change this test to use the SDK, as mocking becomes painful
106-
107-
self.assertEqual(
108-
self.span_attrs,
109-
{"http.status_code": 404, "http.status_text": "NOT FOUND"},
110-
)
111-
112-
def test_internal_error(self):
113-
resp = self.client.get("/hello/500")
114-
self.assertEqual(500, resp.status_code)
86+
resp = self.client.post("/bye")
87+
self.assertEqual(404, resp.status_code)
11588
resp.close()
89+
span_list = self.memory_exporter.get_finished_spans()
90+
self.assertEqual(len(span_list), 1)
91+
self.assertEqual(span_list[0].name, "/bye")
92+
self.assertEqual(span_list[0].kind, trace_api.SpanKind.SERVER)
93+
self.assertEqual(span_list[0].attributes, expected_attrs)
11694

117-
self.start_span.assert_called_with(
118-
"hello_endpoint",
119-
trace_api.INVALID_SPAN_CONTEXT,
120-
kind=trace_api.SpanKind.SERVER,
121-
attributes={
122-
"component": "http",
123-
"http.method": "GET",
124-
"http.server_name": "localhost",
125-
"http.scheme": "http",
126-
"host.port": 80,
127-
"http.host": "localhost",
95+
def test_internal_error(self):
96+
expected_attrs = expected_attributes(
97+
{
12898
"http.target": "/hello/500",
129-
"http.flavor": "1.1",
13099
"http.route": "/hello/<int:helloid>",
131-
},
132-
start_time=mock.ANY,
133-
)
134-
135-
# Nope, this uses Tracer.use_span(end_on_exit)
136-
# self.assertEqual(1, self.span.end.call_count)
137-
# TODO: Change this test to use the SDK, as mocking becomes painful
138-
139-
self.assertEqual(
140-
self.span_attrs,
141-
{
142-
"http.status_code": 500,
143100
"http.status_text": "INTERNAL SERVER ERROR",
144-
},
101+
"http.status_code": 500,
102+
}
145103
)
104+
resp = self.client.get("/hello/500")
105+
self.assertEqual(500, resp.status_code)
106+
resp.close()
107+
span_list = self.memory_exporter.get_finished_spans()
108+
self.assertEqual(len(span_list), 1)
109+
self.assertEqual(span_list[0].name, "hello_endpoint")
110+
self.assertEqual(span_list[0].kind, trace_api.SpanKind.SERVER)
111+
self.assertEqual(span_list[0].attributes, expected_attrs)
146112

147113

148114
if __name__ == "__main__":

‎ext/opentelemetry-ext-testutil/src/opentelemetry/ext/testutil/wsgitestutil.py‎

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,38 @@
11
import io
22
import unittest
3-
import unittest.mock as mock
43
import wsgiref.util as wsgiref_util
4+
from importlib import reload
55

66
from opentelemetry import trace as trace_api
7+
from opentelemetry.sdk.trace import TracerSource, export
8+
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
9+
InMemorySpanExporter,
10+
)
11+
12+
_MEMORY_EXPORTER = None
713

814

915
class WsgiTestBase(unittest.TestCase):
10-
def setUp(self):
11-
self.span = mock.create_autospec(trace_api.Span, spec_set=True)
12-
tracer = trace_api.Tracer()
13-
self.get_tracer_patcher = mock.patch.object(
14-
trace_api.TracerSource,
15-
"get_tracer",
16-
autospec=True,
17-
spec_set=True,
18-
return_value=tracer,
19-
)
20-
self.get_tracer_patcher.start()
21-
22-
self.start_span_patcher = mock.patch.object(
23-
tracer,
24-
"start_span",
25-
autospec=True,
26-
spec_set=True,
27-
return_value=self.span,
16+
@classmethod
17+
def setUpClass(cls):
18+
global _MEMORY_EXPORTER # pylint:disable=global-statement
19+
trace_api.set_preferred_tracer_source_implementation(
20+
lambda T: TracerSource()
2821
)
29-
self.start_span = self.start_span_patcher.start()
22+
tracer_source = trace_api.tracer_source()
23+
_MEMORY_EXPORTER = InMemorySpanExporter()
24+
span_processor = export.SimpleExportSpanProcessor(_MEMORY_EXPORTER)
25+
tracer_source.add_span_processor(span_processor)
26+
27+
@classmethod
28+
def tearDownClass(cls):
29+
reload(trace_api)
30+
31+
def setUp(self):
32+
33+
self.memory_exporter = _MEMORY_EXPORTER
34+
self.memory_exporter.clear()
35+
3036
self.write_buffer = io.BytesIO()
3137
self.write = self.write_buffer.write
3238

@@ -37,10 +43,6 @@ def setUp(self):
3743
self.response_headers = None
3844
self.exc_info = None
3945

40-
def tearDown(self):
41-
self.get_tracer_patcher.stop()
42-
self.start_span_patcher.stop()
43-
4446
def start_response(self, status, response_headers, exc_info=None):
4547
self.status = status
4648
self.response_headers = response_headers

‎ext/opentelemetry-ext-wsgi/tests/test_wsgi_middleware.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -78,10 +78,8 @@ def validate_response(self, response, error=None):
7878
while True:
7979
try:
8080
value = next(response)
81-
self.assertEqual(0, self.span.end.call_count)
8281
self.assertEqual(value, b"*")
8382
except StopIteration:
84-
self.span.end.assert_called_once_with()
8583
break
8684

8785
self.assertEqual(self.status, "200 OK")
@@ -95,12 +93,13 @@ def validate_response(self, response, error=None):
9593
else:
9694
self.assertIsNone(self.exc_info)
9795

98-
# Verify that start_span has been called
99-
self.start_span.assert_called_with(
100-
"/",
101-
trace_api.INVALID_SPAN_CONTEXT,
102-
kind=trace_api.SpanKind.SERVER,
103-
attributes={
96+
span_list = self.memory_exporter.get_finished_spans()
97+
self.assertEqual(len(span_list), 1)
98+
self.assertEqual(span_list[0].name, "/")
99+
self.assertEqual(span_list[0].kind, trace_api.SpanKind.SERVER)
100+
self.assertEqual(
101+
span_list[0].attributes,
102+
{
104103
"component": "http",
105104
"http.method": "GET",
106105
"http.server_name": "127.0.0.1",
@@ -109,6 +108,8 @@ def validate_response(self, response, error=None):
109108
"http.host": "127.0.0.1",
110109
"http.flavor": "1.0",
111110
"http.url": "http://127.0.0.1/",
111+
"http.status_text": "OK",
112+
"http.status_code": 200,
112113
},
113114
)
114115

‎tox.ini‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ commands_pre =
6868
ext: pip install {toxinidir}/opentelemetry-api
6969
wsgi,flask: pip install {toxinidir}/ext/opentelemetry-ext-testutil
7070
wsgi,flask: pip install {toxinidir}/ext/opentelemetry-ext-wsgi
71+
wsgi,flask: pip install {toxinidir}/opentelemetry-sdk
7172
flask: pip install {toxinidir}/ext/opentelemetry-ext-flask[test]
7273
dbapi: pip install {toxinidir}/ext/opentelemetry-ext-dbapi
7374
mysql: pip install {toxinidir}/ext/opentelemetry-ext-dbapi

0 commit comments

Comments
 (0)
⚡