diff --git a/mapping_workbench/backend/task_manager/adapters/task_manager.py b/mapping_workbench/backend/task_manager/adapters/task_manager.py index 262f5f93f..041456744 100644 --- a/mapping_workbench/backend/task_manager/adapters/task_manager.py +++ b/mapping_workbench/backend/task_manager/adapters/task_manager.py @@ -1,3 +1,4 @@ +import socket from concurrent.futures import CancelledError from typing import List @@ -18,7 +19,7 @@ def on_task_done_callback(future): except CancelledError: task.update_task_status(TaskStatus.CANCELED) task.update_exception_message("Task was canceled!") - except TimeoutError as error: + except (socket.timeout, TimeoutError) as error: task.update_task_status(TaskStatus.TIMEOUT) task.update_exception_message(f"Task took longer than {error.args[1]} seconds") except Exception as error: diff --git a/tests/unit/backend/task_manager/test_task_manager.py b/tests/unit/backend/task_manager/test_task_manager.py index b78d914d0..a59494288 100644 --- a/tests/unit/backend/task_manager/test_task_manager.py +++ b/tests/unit/backend/task_manager/test_task_manager.py @@ -44,7 +44,7 @@ def test_task_manager(): while task_manager.is_active(): time.sleep(1) - assert task_manager.tasks[0].task_metadata.exception_message == "Task took longer than 1 seconds" + assert task_manager.tasks[0].task_metadata.exception_message == "Task took longer than 1.0 seconds" assert task_manager.tasks[0].get_task_status() == TaskStatus.TIMEOUT task_manager.delete_all_tasks()