1import pytest
2import six
3from webtest import TestApp
4
5from ddtrace import config
6from ddtrace.contrib.wsgi import wsgi
7from ddtrace.internal.compat import PY2
8from ddtrace.internal.compat import PY3
9
10
11if PY2:
12    import exceptions
13
14    generatorExit = exceptions.GeneratorExit
15else:
16    import builtins
17
18    generatorExit = builtins.GeneratorExit
19
20from tests.utils import override_config
21from tests.utils import override_http_config
22from tests.utils import snapshot
23
24
25def chunked_response(start_response):
26    status = "200 OK"
27    headers = [("Content-type", "text/plain")]
28    start_response(status, headers)
29    for i in range(1000):
30        yield b"%d" % i
31
32
33def chunked_response_generator_error(start_response):
34    status = "200 OK"
35    headers = [("Content-type", "text/plain")]
36    start_response(status, headers)
37    for i in range(1000):
38        if i < 999:
39            yield b"%d" % i
40        else:
41            raise generatorExit()
42
43
44def application(environ, start_response):
45    if environ["PATH_INFO"] == "/error":
46        raise Exception("Oops!")
47    elif environ["PATH_INFO"] == "/chunked":
48        return chunked_response(start_response)
49    elif environ["PATH_INFO"] == "/generatorError":
50        return chunked_response_generator_error(start_response)
51    else:
52        body = six.b("<html><body><h1>Hello World</h1></body></html>")
53        headers = [
54            ("Content-Type", "text/html; charset=utf8"),
55            ("Content-Length", str(len(body))),
56            ("my-response-header", "test_response_value"),
57        ]
58        start_response("200 OK", headers)
59        return [body]
60
61
62def test_middleware(tracer, test_spans):
63    app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
64    resp = app.get("/")
65    assert resp.status == "200 OK"
66    assert resp.status_int == 200
67    spans = test_spans.pop()
68    assert len(spans) == 4
69
70    with pytest.raises(Exception):
71        app.get("/error")
72
73    spans = test_spans.pop()
74    assert len(spans) == 2
75    assert spans[0].error == 1
76
77
78def test_distributed_tracing(tracer, test_spans):
79    app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
80    resp = app.get("/", headers={"X-Datadog-Parent-Id": "1234", "X-Datadog-Trace-Id": "4321"})
81
82    assert config.wsgi.distributed_tracing is True
83    assert resp.status == "200 OK"
84    assert resp.status_int == 200
85
86    spans = test_spans.pop()
87    assert len(spans) == 4
88    root = spans[0]
89    assert root.name == "wsgi.request"
90    assert root.trace_id == 4321
91    assert root.parent_id == 1234
92
93    with override_config("wsgi", dict(distributed_tracing=False)):
94        app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
95        resp = app.get("/", headers={"X-Datadog-Parent-Id": "1234", "X-Datadog-Trace-Id": "4321"})
96        assert config.wsgi.distributed_tracing is False
97        assert resp.status == "200 OK"
98        assert resp.status_int == 200
99
100        spans = test_spans.pop()
101        assert len(spans) == 4
102        root = spans[0]
103        assert root.name == "wsgi.request"
104        assert root.trace_id != 4321
105        assert root.parent_id != 1234
106
107
108def test_query_string_tracing(tracer, test_spans):
109    with override_http_config("wsgi", dict(trace_query_string=True)):
110        app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
111        response = app.get("/?foo=bar&x=y")
112
113        assert response.status_int == 200
114        assert response.status == "200 OK"
115
116        spans = test_spans.pop_traces()
117        assert len(spans) == 1
118        assert len(spans[0]) == 4
119        request_span = spans[0][0]
120        assert request_span.service == "wsgi"
121        assert request_span.name == "wsgi.request"
122        assert request_span.resource == "GET /"
123        assert request_span.error == 0
124        assert request_span.get_tag("http.method") == "GET"
125        assert request_span.get_tag("http.status_code") == "200"
126        assert request_span.get_tag("http.query.string") == "foo=bar&x=y"
127
128        assert spans[0][1].name == "wsgi.application"
129        assert spans[0][2].name == "wsgi.start_response"
130        assert spans[0][3].name == "wsgi.response"
131
132
133def test_http_request_header_tracing(tracer, test_spans):
134    config.wsgi.http.trace_headers(["my-header"])
135    app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
136    resp = app.get("/", headers={"my-header": "test_value"})
137
138    assert resp.status == "200 OK"
139    assert resp.status_int == 200
140    spans = test_spans.pop()
141    assert spans[0].get_tag("http.request.headers.my-header") == "test_value"
142
143
144def test_http_response_header_tracing(tracer, test_spans):
145    config.wsgi.http.trace_headers(["my-response-header"])
146    app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
147    resp = app.get("/", headers={"my-header": "test_value"})
148
149    assert resp.status == "200 OK"
150    assert resp.status_int == 200
151
152    spans = test_spans.pop()
153    assert spans[0].get_tag("http.response.headers.my-response-header") == "test_response_value"
154
155
156def test_service_name_can_be_overriden(tracer, test_spans):
157    with override_config("wsgi", dict(service_name="test-override-service")):
158        app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
159        response = app.get("/")
160        assert response.status_code == 200
161
162        spans = test_spans.pop_traces()
163        assert len(spans) > 0
164        span = spans[0][0]
165        assert span.service == "test-override-service"
166
167
168def test_generator_exit_ignored_in_top_level_span(tracer, test_spans):
169    with pytest.raises(generatorExit):
170        app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
171        app.get("/generatorError")
172
173    spans = test_spans.pop()
174    assert spans[2].error == 1
175    assert "GeneratorExit" in spans[2].get_tag("error.type")
176    assert spans[0].error == 0
177
178
179@snapshot(ignores=["meta.error.stack"], variants={"py2": PY2, "py3": PY3})
180def test_generator_exit_ignored_in_top_level_span_snapshot():
181    with pytest.raises(generatorExit):
182        app = TestApp(wsgi.DDWSGIMiddleware(application))
183        app.get("/generatorError")
184
185
186def test_chunked_response(tracer, test_spans):
187    app = TestApp(wsgi.DDWSGIMiddleware(application, tracer=tracer))
188    resp = app.get("/chunked")
189    assert resp.status == "200 OK"
190    assert resp.status_int == 200
191    assert resp.text.startswith("0123456789")
192    assert resp.text.endswith("999")
193
194    spans = test_spans.pop_traces()
195    span = spans[0][0]
196    assert span.resource == "GET /chunked"
197    assert span.name == "wsgi.request"
198
199
200@snapshot()
201def test_chunked():
202    app = TestApp(wsgi.DDWSGIMiddleware(application))
203    resp = app.get("/chunked")
204    assert resp.status == "200 OK"
205    assert resp.status_int == 200
206    assert resp.text.startswith("0123456789")
207    assert resp.text.endswith("999")
208
209
210@snapshot()
211def test_200():
212    app = TestApp(wsgi.DDWSGIMiddleware(application))
213    resp = app.get("/")
214    assert resp.status == "200 OK"
215    assert resp.status_int == 200
216
217
218@snapshot(ignores=["meta.error.stack"], variants={"py2": PY2, "py3": PY3})
219def test_500():
220    app = TestApp(wsgi.DDWSGIMiddleware(application))
221    with pytest.raises(Exception):
222        app.get("/error")
223