16
16
17
17
import asyncio
18
18
import sys
19
+ import time
19
20
import unittest
20
21
from timeit import default_timer
21
22
from unittest import mock
57
58
"http.server.request.size" : _duration_attrs ,
58
59
}
59
60
61
+ _SIMULATED_BACKGROUND_TASK_EXECUTION_TIME_S = 0.01
62
+
60
63
61
64
async def http_app (scope , receive , send ):
62
65
message = await receive ()
@@ -99,6 +102,108 @@ async def simple_asgi(scope, receive, send):
99
102
await websocket_app (scope , receive , send )
100
103
101
104
105
+ async def long_response_asgi (scope , receive , send ):
106
+ assert isinstance (scope , dict )
107
+ assert scope ["type" ] == "http"
108
+ message = await receive ()
109
+ scope ["headers" ] = [(b"content-length" , b"128" )]
110
+ assert scope ["type" ] == "http"
111
+ if message .get ("type" ) == "http.request" :
112
+ await send (
113
+ {
114
+ "type" : "http.response.start" ,
115
+ "status" : 200 ,
116
+ "headers" : [
117
+ [b"Content-Type" , b"text/plain" ],
118
+ [b"content-length" , b"1024" ],
119
+ ],
120
+ }
121
+ )
122
+ await send (
123
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : True }
124
+ )
125
+ await send (
126
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : True }
127
+ )
128
+ await send (
129
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : True }
130
+ )
131
+ await send (
132
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : False }
133
+ )
134
+
135
+
136
+ async def background_execution_asgi (scope , receive , send ):
137
+ assert isinstance (scope , dict )
138
+ assert scope ["type" ] == "http"
139
+ message = await receive ()
140
+ scope ["headers" ] = [(b"content-length" , b"128" )]
141
+ assert scope ["type" ] == "http"
142
+ if message .get ("type" ) == "http.request" :
143
+ await send (
144
+ {
145
+ "type" : "http.response.start" ,
146
+ "status" : 200 ,
147
+ "headers" : [
148
+ [b"Content-Type" , b"text/plain" ],
149
+ [b"content-length" , b"1024" ],
150
+ ],
151
+ }
152
+ )
153
+ await send (
154
+ {
155
+ "type" : "http.response.body" ,
156
+ "body" : b"*" ,
157
+ }
158
+ )
159
+ time .sleep (_SIMULATED_BACKGROUND_TASK_EXECUTION_TIME_S )
160
+
161
+
162
+ async def background_execution_trailers_asgi (scope , receive , send ):
163
+ assert isinstance (scope , dict )
164
+ assert scope ["type" ] == "http"
165
+ message = await receive ()
166
+ scope ["headers" ] = [(b"content-length" , b"128" )]
167
+ assert scope ["type" ] == "http"
168
+ if message .get ("type" ) == "http.request" :
169
+ await send (
170
+ {
171
+ "type" : "http.response.start" ,
172
+ "status" : 200 ,
173
+ "headers" : [
174
+ [b"Content-Type" , b"text/plain" ],
175
+ [b"content-length" , b"1024" ],
176
+ ],
177
+ "trailers" : True ,
178
+ }
179
+ )
180
+ await send (
181
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : True }
182
+ )
183
+ await send (
184
+ {"type" : "http.response.body" , "body" : b"*" , "more_body" : False }
185
+ )
186
+ await send (
187
+ {
188
+ "type" : "http.response.trailers" ,
189
+ "headers" : [
190
+ [b"trailer" , b"test-trailer" ],
191
+ ],
192
+ "more_trailers" : True ,
193
+ }
194
+ )
195
+ await send (
196
+ {
197
+ "type" : "http.response.trailers" ,
198
+ "headers" : [
199
+ [b"trailer" , b"second-test-trailer" ],
200
+ ],
201
+ "more_trailers" : False ,
202
+ }
203
+ )
204
+ time .sleep (_SIMULATED_BACKGROUND_TASK_EXECUTION_TIME_S )
205
+
206
+
102
207
async def error_asgi (scope , receive , send ):
103
208
assert isinstance (scope , dict )
104
209
assert scope ["type" ] == "http"
@@ -127,14 +232,19 @@ def validate_outputs(self, outputs, error=None, modifiers=None):
127
232
# Ensure modifiers is a list
128
233
modifiers = modifiers or []
129
234
# Check for expected outputs
130
- self .assertEqual (len (outputs ), 2 )
131
235
response_start = outputs [0 ]
132
- response_body = outputs [1 ]
236
+ response_final_body = [
237
+ output
238
+ for output in outputs
239
+ if output ["type" ] == "http.response.body"
240
+ ][- 1 ]
241
+
133
242
self .assertEqual (response_start ["type" ], "http.response.start" )
134
- self .assertEqual (response_body ["type" ], "http.response.body" )
243
+ self .assertEqual (response_final_body ["type" ], "http.response.body" )
244
+ self .assertEqual (response_final_body .get ("more_body" , False ), False )
135
245
136
246
# Check http response body
137
- self .assertEqual (response_body ["body" ], b"*" )
247
+ self .assertEqual (response_final_body ["body" ], b"*" )
138
248
139
249
# Check http response start
140
250
self .assertEqual (response_start ["status" ], 200 )
@@ -153,7 +263,6 @@ def validate_outputs(self, outputs, error=None, modifiers=None):
153
263
154
264
# Check spans
155
265
span_list = self .memory_exporter .get_finished_spans ()
156
- self .assertEqual (len (span_list ), 4 )
157
266
expected = [
158
267
{
159
268
"name" : "GET / http receive" ,
@@ -194,6 +303,7 @@ def validate_outputs(self, outputs, error=None, modifiers=None):
194
303
for modifier in modifiers :
195
304
expected = modifier (expected )
196
305
# Check that output matches
306
+ self .assertEqual (len (span_list ), len (expected ))
197
307
for span , expected in zip (span_list , expected ):
198
308
self .assertEqual (span .name , expected ["name" ])
199
309
self .assertEqual (span .kind , expected ["kind" ])
@@ -232,6 +342,80 @@ def test_asgi_exc_info(self):
232
342
outputs = self .get_all_output ()
233
343
self .validate_outputs (outputs , error = ValueError )
234
344
345
+ def test_long_response (self ):
346
+ """Test that the server span is ended on the final response body message.
347
+
348
+ If the server span is ended early then this test will fail due
349
+ to discrepancies in the expected list of spans and the emitted list of spans.
350
+ """
351
+ app = otel_asgi .OpenTelemetryMiddleware (long_response_asgi )
352
+ self .seed_app (app )
353
+ self .send_default_request ()
354
+ outputs = self .get_all_output ()
355
+
356
+ def add_more_body_spans (expected : list ):
357
+ more_body_span = {
358
+ "name" : "GET / http send" ,
359
+ "kind" : trace_api .SpanKind .INTERNAL ,
360
+ "attributes" : {"type" : "http.response.body" },
361
+ }
362
+ extra_spans = [more_body_span ] * 3
363
+ expected [2 :2 ] = extra_spans
364
+ return expected
365
+
366
+ self .validate_outputs (outputs , modifiers = [add_more_body_spans ])
367
+
368
+ def test_background_execution (self ):
369
+ """Test that the server span is ended BEFORE the background task is finished."""
370
+ app = otel_asgi .OpenTelemetryMiddleware (background_execution_asgi )
371
+ self .seed_app (app )
372
+ self .send_default_request ()
373
+ outputs = self .get_all_output ()
374
+ self .validate_outputs (outputs )
375
+ span_list = self .memory_exporter .get_finished_spans ()
376
+ server_span = span_list [- 1 ]
377
+ assert server_span .kind == SpanKind .SERVER
378
+ span_duration_nanos = server_span .end_time - server_span .start_time
379
+ self .assertLessEqual (
380
+ span_duration_nanos ,
381
+ _SIMULATED_BACKGROUND_TASK_EXECUTION_TIME_S * 10 ** 9 ,
382
+ )
383
+
384
+ def test_trailers (self ):
385
+ """Test that trailers are emitted as expected and that the server span is ended
386
+ BEFORE the background task is finished."""
387
+ app = otel_asgi .OpenTelemetryMiddleware (
388
+ background_execution_trailers_asgi
389
+ )
390
+ self .seed_app (app )
391
+ self .send_default_request ()
392
+ outputs = self .get_all_output ()
393
+
394
+ def add_body_and_trailer_span (expected : list ):
395
+ body_span = {
396
+ "name" : "GET / http send" ,
397
+ "kind" : trace_api .SpanKind .INTERNAL ,
398
+ "attributes" : {"type" : "http.response.body" },
399
+ }
400
+ trailer_span = {
401
+ "name" : "GET / http send" ,
402
+ "kind" : trace_api .SpanKind .INTERNAL ,
403
+ "attributes" : {"type" : "http.response.trailers" },
404
+ }
405
+ expected [2 :2 ] = [body_span ]
406
+ expected [4 :4 ] = [trailer_span ] * 2
407
+ return expected
408
+
409
+ self .validate_outputs (outputs , modifiers = [add_body_and_trailer_span ])
410
+ span_list = self .memory_exporter .get_finished_spans ()
411
+ server_span = span_list [- 1 ]
412
+ assert server_span .kind == SpanKind .SERVER
413
+ span_duration_nanos = server_span .end_time - server_span .start_time
414
+ self .assertLessEqual (
415
+ span_duration_nanos ,
416
+ _SIMULATED_BACKGROUND_TASK_EXECUTION_TIME_S * 10 ** 9 ,
417
+ )
418
+
235
419
def test_override_span_name (self ):
236
420
"""Test that default span_names can be overwritten by our callback function."""
237
421
span_name = "Dymaxion"
0 commit comments