Skip to content

Commit

Permalink
Fix websocket span messages (#426)
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexmojaki
alexmojaki authored Sep 16, 2024
1 parent f2093b3 commit 965904d
Show file tree
Hide file tree
Showing 3 changed files with 171 additions and 18 deletions.
3 changes: 3 additions & 0 deletions logfire/_internal/exporters/processor_wrapper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,9 @@ def _tweak_http_spans(span: ReadableSpanDict):
except Exception: # pragma: no cover
pass

if not method and name in ('HTTP', f'HTTP {target}', f'HTTP {route}'):
method = 'HTTP'

# Build up a list of possible span names and messages in order from worst to best
names: list[str] = []
messages: list[str] = []
Expand Down
150 changes: 150 additions & 0 deletions tests/otel_integrations/test_fastapi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,14 @@ async def echo_body(request: Request):
return await request.body()


async def websocket_endpoint(websocket: WebSocket, name: str):
logfire.info('websocket_endpoint: {name}', name=name)
await websocket.accept()
assert (await websocket.receive_text()) == 'ping'
await websocket.send_text('pong')
await websocket.close()


@pytest.fixture()
def app():
# Don't define the endpoint functions in this fixture to prevent a qualname with <locals> in it
Expand All @@ -90,6 +98,7 @@ def app():
app.get('/validation_error')(validation_error)
app.get('/with_path_param/{param}')(with_path_param)
app.get('/secret/{path_param}', name='secret')(get_secret)
app.websocket('/ws/{name}')(websocket_endpoint)
first_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
second_lvl_app.get('/other', name='other_route_name', operation_id='other_route_operation_id')(other_route)
return app
Expand Down Expand Up @@ -1818,3 +1827,144 @@ def test_request_hooks_with_send_receive_spans(exporter: TestExporter):
},
]
)


def test_websocket(client: TestClient, exporter: TestExporter) -> None:
with client.websocket_connect('/ws/foo') as websocket:
websocket.send_text('ping')
data = websocket.receive_text()
assert data == 'pong'

assert exporter.exported_spans_as_dict() == snapshot(
[
{
'name': 'FastAPI arguments',
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 2000000000,
'end_time': 3000000000,
'attributes': {
'logfire.msg_template': 'FastAPI arguments',
'logfire.msg': 'FastAPI arguments',
'logfire.span_type': 'span',
'logfire.level_num': 5,
},
},
{
'name': 'websocket_endpoint: {name}',
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 4000000000,
'end_time': 4000000000,
'attributes': {
'logfire.span_type': 'log',
'logfire.level_num': 9,
'logfire.msg_template': 'websocket_endpoint: {name}',
'logfire.msg': 'websocket_endpoint: foo',
'code.filepath': 'test_fastapi.py',
'code.function': 'websocket_endpoint',
'code.lineno': 123,
'name': 'foo',
'logfire.json_schema': '{"type":"object","properties":{"name":{}}}',
},
},
{
'name': 'HTTP /ws/{name} websocket receive connect',
'context': {'trace_id': 1, 'span_id': 6, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 5000000000,
'end_time': 6000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket receive connect',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.connect',
},
},
{
'name': 'HTTP /ws/{name} websocket send accept',
'context': {'trace_id': 1, 'span_id': 8, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 7000000000,
'end_time': 8000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send accept',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.accept',
},
},
{
'name': 'HTTP /ws/{name} websocket receive',
'context': {'trace_id': 1, 'span_id': 10, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 9000000000,
'end_time': 10000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket receive',
'logfire.level_num': 5,
'http.status_code': 200,
'http.response.status_code': 200,
'asgi.event.type': 'websocket.receive',
},
},
{
'name': 'HTTP /ws/{name} websocket send',
'context': {'trace_id': 1, 'span_id': 12, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 11000000000,
'end_time': 12000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.send',
'http.status_code': 200,
'http.response.status_code': 200,
},
},
{
'name': 'HTTP /ws/{name} websocket send close',
'context': {'trace_id': 1, 'span_id': 14, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 13000000000,
'end_time': 14000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/{name} websocket send close',
'logfire.level_num': 5,
'asgi.event.type': 'websocket.close',
},
},
{
'name': 'HTTP /ws/{name}',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 15000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': 'HTTP /ws/foo',
'http.scheme': 'ws',
'url.scheme': 'ws',
'http.host': 'testserver',
'client.address': 'testserver',
'net.host.port': 80,
'server.port': 80,
'http.target': '/ws/foo',
'url.path': '/ws/foo',
'http.url': 'ws://testserver/ws/foo',
'http.server_name': 'testserver',
'http.user_agent': 'testclient',
'user_agent.original': 'testclient',
'net.peer.ip': 'testclient',
'net.peer.port': 50000,
'client.port': 50000,
'http.route': '/ws/{name}',
'http.status_code': 200,
'http.response.status_code': 200,
},
},
]
)
36 changes: 18 additions & 18 deletions tests/otel_integrations/test_starlette.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ async def websocket_endpoint(websocket: WebSocket):
def app():
routes = [
Route('/secret/{path_param}', secret),
WebSocketRoute('/ws', websocket_endpoint),
WebSocketRoute('/ws/{name}', websocket_endpoint),
]

app = Starlette(routes=routes)
Expand All @@ -47,107 +47,107 @@ def client(app: Starlette) -> TestClient:


def test_websocket(client: TestClient, exporter: TestExporter) -> None:
with client.websocket_connect('/ws') as websocket:
with client.websocket_connect('/ws/foo') as websocket:
websocket.send_text('ping')
data = websocket.receive_text()
assert data == 'pong'

assert exporter.exported_spans_as_dict() == snapshot(
[
{
'name': '/ws websocket receive connect',
'name': '/ws/{name} websocket receive connect',
'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 2000000000,
'end_time': 3000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket receive connect',
'logfire.msg': '/ws/{name} websocket receive connect',
'asgi.event.type': 'websocket.connect',
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket send accept',
'name': '/ws/{name} websocket send accept',
'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 4000000000,
'end_time': 5000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send accept',
'logfire.msg': '/ws/{name} websocket send accept',
'asgi.event.type': 'websocket.accept',
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket receive',
'name': '/ws/{name} websocket receive',
'context': {'trace_id': 1, 'span_id': 7, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 6000000000,
'end_time': 7000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket receive',
'logfire.msg': '/ws/{name} websocket receive',
'http.status_code': 200,
'asgi.event.type': 'websocket.receive',
'http.response.status_code': 200,
'logfire.level_num': 5,
},
},
{
'name': '/ws websocket send',
'name': '/ws/{name} websocket send',
'context': {'trace_id': 1, 'span_id': 9, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 8000000000,
'end_time': 9000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send',
'logfire.msg': '/ws/{name} websocket send',
'http.status_code': 200,
'asgi.event.type': 'websocket.send',
'logfire.level_num': 5,
'http.response.status_code': 200,
},
},
{
'name': '/ws websocket send close',
'name': '/ws/{name} websocket send close',
'context': {'trace_id': 1, 'span_id': 11, 'is_remote': False},
'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'start_time': 10000000000,
'end_time': 11000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws websocket send close',
'logfire.msg': '/ws/{name} websocket send close',
'asgi.event.type': 'websocket.close',
'logfire.level_num': 5,
},
},
{
'name': '/ws',
'name': '/ws/{name}',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': 1000000000,
'end_time': 12000000000,
'attributes': {
'logfire.span_type': 'span',
'logfire.msg': '/ws',
'logfire.msg': '/ws/foo',
'http.scheme': 'ws',
'url.scheme': 'ws',
'http.host': 'testserver',
'net.host.port': 80,
'server.port': 80,
'http.target': '/ws',
'url.path': '/ws',
'http.url': 'ws://testserver/ws',
'http.target': '/ws/foo',
'url.path': '/ws/foo',
'http.url': 'ws://testserver/ws/foo',
'http.server_name': 'testserver',
'http.user_agent': 'testclient',
'user_agent.original': 'testclient',
'net.peer.ip': 'testclient',
'client.address': 'testserver',
'net.peer.port': 50000,
'client.port': 50000,
'http.route': '/ws',
'http.route': '/ws/{name}',
'http.request.header.host': ('testserver',),
'http.request.header.accept': ('*/*',),
'http.request.header.accept_encoding': ('gzip, deflate',),
Expand Down

0 comments on commit 965904d

Please sign in to comment.