Skip to content

Commit

Permalink
Enhance robustness
Browse files Browse the repository at this point in the history
  • Loading branch information
yihong1120 committed Dec 17, 2024
1 parent e8711fb commit cc23d34
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 6 deletions.
3 changes: 2 additions & 1 deletion examples/streaming_web/backend/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -362,7 +362,8 @@ def verify_localhost(request: Request) -> None:
Raises:
HTTPException: If the request is not made from localhost.
"""
if request.client.host not in ['127.0.0.1', '::1']:
client = request.client
if client is None or client.host not in ['127.0.0.1', '::1']:
raise HTTPException(
status_code=403,
detail='Access is restricted to localhost only.',
Expand Down
13 changes: 8 additions & 5 deletions tests/examples/streaming_web/backend/app_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,25 +101,28 @@ def test_app_running_configuration(
)

@patch(
"examples.streaming_web.backend.app.redis_manager.client",
'examples.streaming_web.backend.app.redis_manager.client',
new_callable=AsyncMock,
)
@patch(
"examples.streaming_web.backend.app.FastAPILimiter.init",
'examples.streaming_web.backend.app.FastAPILimiter.init',
new_callable=AsyncMock,
)
def test_lifespan_events(self, mock_limiter_init: AsyncMock, mock_redis_client: AsyncMock) -> None:
def test_lifespan_events(
self,
mock_limiter_init: AsyncMock,
mock_redis_client: AsyncMock,
) -> None:
"""
Test that the lifespan events are properly handled.
"""
with TestClient(app_module.app) as client:
response = client.get("/")
response = client.get('/')
# 驗證返回的狀態碼是 404
self.assertEqual(response.status_code, 404)
mock_limiter_init.assert_called_once_with(mock_redis_client)
mock_redis_client.close.assert_called_once()


def tearDown(self) -> None:
del self.client

Expand Down

0 comments on commit cc23d34

Please sign in to comment.