Skip to content

Commit 7c966a0

Browse files
tech4242tech4242
committed
fix: linter
1 parent 2b98f56 commit 7c966a0

File tree

2 files changed

+90
-98
lines changed

2 files changed

+90
-98
lines changed

‎tests/test_broadcaster.py‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
import asyncio
2-
from unittest.mock import AsyncMock, MagicMock
1+
from unittest.mock import AsyncMock
32

43
import pytest
54

@@ -11,25 +10,25 @@ async def test_broadcast_new_log_with_clients():
1110
# Create mock websocket clients
1211
mock_client1 = AsyncMock()
1312
mock_client2 = AsyncMock()
14-
13+
1514
# Add to active clients
1615
active_clients.append(mock_client1)
1716
active_clients.append(mock_client2)
18-
17+
1918
try:
2019
# Broadcast a message
2120
test_log = {
2221
"timestamp": "2024-01-01T12:00:00Z",
2322
"src_ip": "127.0.0.1",
2423
"message": "test"
2524
}
26-
25+
2726
await broadcast_new_log(test_log)
28-
27+
2928
# Verify both clients received the message
3029
mock_client1.send_json.assert_called_once_with(test_log)
3130
mock_client2.send_json.assert_called_once_with(test_log)
32-
31+
3332
finally:
3433
# Clean up
3534
active_clients.clear()
@@ -40,36 +39,36 @@ async def test_broadcast_new_log_with_disconnected_client():
4039
# Create mock websocket clients
4140
mock_client1 = AsyncMock()
4241
mock_client2 = AsyncMock()
43-
42+
4443
# Make client1 throw an exception (disconnected)
4544
mock_client1.send_json.side_effect = Exception("Disconnected")
46-
45+
4746
# Add to active clients
4847
active_clients.append(mock_client1)
4948
active_clients.append(mock_client2)
50-
49+
5150
try:
5251
# Broadcast a message
5352
test_log = {"test": "data"}
54-
53+
5554
await broadcast_new_log(test_log)
56-
55+
5756
# Client1 should be removed from active_clients
5857
assert mock_client1 not in active_clients
5958
assert mock_client2 in active_clients
60-
59+
6160
# Client2 should still receive the message
6261
mock_client2.send_json.assert_called_once_with(test_log)
63-
62+
6463
finally:
6564
# Clean up
6665
active_clients.clear()
6766

6867

69-
@pytest.mark.asyncio
68+
@pytest.mark.asyncio
7069
async def test_broadcast_new_log_no_clients():
7170
# Ensure list is empty
7271
active_clients.clear()
73-
72+
7473
# Should not raise any errors
75-
await broadcast_new_log({"test": "data"})
74+
await broadcast_new_log({"test": "data"})

‎tests/test_web_server.py‎

Lines changed: 74 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import json
21
from datetime import datetime, timezone
32
from unittest.mock import MagicMock, patch
43

@@ -62,125 +61,119 @@ def test_logs_endpoint_with_limit(client, mock_fetch_logs):
6261

6362

6463
def test_websocket_connection():
65-
with TestClient(app) as client:
66-
with client.websocket_connect("/ws") as websocket:
67-
# WebSocket should connect successfully
68-
# The connection is automatically closed when exiting the context
69-
pass
64+
with TestClient(app) as client, client.websocket_connect("/ws"):
65+
# WebSocket should connect successfully
66+
# The connection is automatically closed when exiting the context
67+
pass
7068

7169

7270
def test_websocket_ping_pong():
73-
import time
7471
import asyncio
75-
with TestClient(app) as client:
72+
with TestClient(app) as client, patch("asyncio.wait_for") as mock_wait_for:
7673
# Mock asyncio.wait_for to simulate timeout
77-
with patch("asyncio.wait_for") as mock_wait_for:
78-
mock_wait_for.side_effect = asyncio.TimeoutError()
79-
80-
with client.websocket_connect("/ws") as websocket:
81-
# This should trigger the timeout handling
82-
try:
83-
# The server should send a ping
84-
data = websocket.receive_json()
85-
assert data == {"type": "ping"}
86-
except Exception:
87-
# Connection might close, that's ok
88-
pass
74+
mock_wait_for.side_effect = asyncio.TimeoutError()
75+
76+
with client.websocket_connect("/ws") as websocket:
77+
# This should trigger the timeout handling
78+
try:
79+
# The server should send a ping
80+
data = websocket.receive_json()
81+
assert data == {"type": "ping"}
82+
except Exception:
83+
# Connection might close, that's ok
84+
pass
8985

9086

9187
def test_websocket_disconnect():
92-
with TestClient(app) as client:
93-
with client.websocket_connect("/ws") as websocket:
94-
# Force disconnect by closing
95-
websocket.close()
96-
# Should handle gracefully
88+
with TestClient(app) as client, client.websocket_connect("/ws") as websocket:
89+
# Force disconnect by closing
90+
websocket.close()
91+
# Should handle gracefully
9792

9893

9994
def test_websocket_broadcast():
100-
with TestClient(app) as client:
95+
with TestClient(app) as client, client.websocket_connect("/ws") as websocket:
10196
# Test that websocket connects and can receive messages
102-
with client.websocket_connect("/ws") as websocket:
103-
# Just verify the websocket connection works
104-
assert websocket is not None
105-
# Could send a test message if needed
106-
# websocket.send_json({"type": "ping"})
97+
# Just verify the websocket connection works
98+
assert websocket is not None
99+
# Could send a test message if needed
100+
# websocket.send_json({"type": "ping"})
107101

108102

109103

110104

111105
def test_start_sniffer_thread():
112106
from mcphawk.web.server import _start_sniffer_thread
113-
114-
with patch("mcphawk.sniffer.start_sniffer") as mock_sniffer:
115-
with patch("threading.Thread") as mock_thread:
107+
108+
with patch("mcphawk.sniffer.start_sniffer"), patch("threading.Thread") as mock_thread:
116109
mock_thread_instance = MagicMock()
117110
mock_thread.return_value = mock_thread_instance
118-
111+
119112
_start_sniffer_thread("tcp port 3000", auto_detect=False, debug=True)
120-
113+
121114
# Verify thread was created with correct arguments
122115
mock_thread.assert_called_once()
123116
args, kwargs = mock_thread.call_args
124117
assert kwargs["daemon"] is True
125118
assert callable(kwargs["target"])
126-
119+
127120
# Verify thread was started
128121
mock_thread_instance.start.assert_called_once()
129122

130123

131124
def test_run_web_with_sniffer():
132125
from mcphawk.web.server import run_web
133-
134-
with patch("uvicorn.run") as mock_uvicorn:
135-
with patch("mcphawk.web.server._start_sniffer_thread") as mock_start_sniffer:
136-
with patch("mcphawk.logger.init_db") as mock_init_db:
137-
run_web(
138-
sniffer=True,
139-
host="0.0.0.0",
140-
port=9000,
141-
filter_expr="tcp port 3000",
142-
auto_detect=False,
143-
debug=True
144-
)
145-
146-
# Verify sniffer was started
147-
mock_start_sniffer.assert_called_once_with("tcp port 3000", False, True)
148-
149-
# Verify uvicorn was started with correct params
150-
mock_uvicorn.assert_called_once()
151-
args, kwargs = mock_uvicorn.call_args
152-
assert args[0] == app # First arg is the app
153-
assert kwargs["host"] == "0.0.0.0"
154-
assert kwargs["port"] == 9000
126+
127+
with patch("uvicorn.run") as mock_uvicorn, \
128+
patch("mcphawk.web.server._start_sniffer_thread") as mock_start_sniffer, \
129+
patch("mcphawk.logger.init_db"):
130+
run_web(
131+
sniffer=True,
132+
host="0.0.0.0",
133+
port=9000,
134+
filter_expr="tcp port 3000",
135+
auto_detect=False,
136+
debug=True
137+
)
138+
139+
# Verify sniffer was started
140+
mock_start_sniffer.assert_called_once_with("tcp port 3000", False, True)
141+
142+
# Verify uvicorn was started with correct params
143+
mock_uvicorn.assert_called_once()
144+
args, kwargs = mock_uvicorn.call_args
145+
assert args[0] == app # First arg is the app
146+
assert kwargs["host"] == "0.0.0.0"
147+
assert kwargs["port"] == 9000
155148

156149

157150
def test_run_web_without_sniffer():
158151
from mcphawk.web.server import run_web
159-
160-
with patch("uvicorn.run") as mock_uvicorn:
161-
with patch("mcphawk.web.server._start_sniffer_thread") as mock_start_sniffer:
162-
with patch("mcphawk.logger.init_db") as mock_init_db:
163-
run_web(
164-
sniffer=False,
165-
host="127.0.0.1",
166-
port=8000,
167-
debug=False
168-
)
169-
170-
# Verify sniffer was NOT started
171-
mock_start_sniffer.assert_not_called()
172-
173-
# Verify uvicorn was started with correct params
174-
mock_uvicorn.assert_called_once()
175-
args, kwargs = mock_uvicorn.call_args
176-
assert args[0] == app # First arg is the app
177-
assert kwargs["host"] == "127.0.0.1"
178-
assert kwargs["port"] == 8000
152+
153+
with patch("uvicorn.run") as mock_uvicorn, \
154+
patch("mcphawk.web.server._start_sniffer_thread") as mock_start_sniffer, \
155+
patch("mcphawk.logger.init_db"):
156+
run_web(
157+
sniffer=False,
158+
host="127.0.0.1",
159+
port=8000,
160+
debug=False
161+
)
162+
163+
# Verify sniffer was NOT started
164+
mock_start_sniffer.assert_not_called()
165+
166+
# Verify uvicorn was started with correct params
167+
mock_uvicorn.assert_called_once()
168+
args, kwargs = mock_uvicorn.call_args
169+
assert args[0] == app # First arg is the app
170+
assert kwargs["host"] == "127.0.0.1"
171+
assert kwargs["port"] == 8000
179172

180173

181174
def test_run_web_sniffer_without_filter():
182175
from mcphawk.web.server import run_web
183-
176+
184177
# Test that ValueError is raised when sniffer=True but no filter_expr
185178
with pytest.raises(ValueError, match="filter_expr is required"):
186-
run_web(sniffer=True, filter_expr=None)
179+
run_web(sniffer=True, filter_expr=None)

0 commit comments

Comments
 (0)