Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion internal_filesystem/lib/mpos/task_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,11 +48,20 @@ def enable(cls):
def disable(cls):
cls.disabled = True

@classmethod
async def _run_task(cls, coroutine):
task = asyncio.current_task()
try:
return await coroutine
finally:
if task in cls.task_list:
cls.task_list.remove(task)

@classmethod
def create_task(cls, coroutine):
if cls.disabled:
return None
task = asyncio.create_task(coroutine)
task = asyncio.create_task(cls._run_task(coroutine))
cls.task_list.append(task)
return task

Expand Down
103 changes: 103 additions & 0 deletions tests/test_task_manager.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,18 @@
import asyncio
import unittest
from mpos.task_manager import TaskManager


class TestTaskManagerState(unittest.TestCase):

def setUp(self):
self.original_task_list = TaskManager.task_list
TaskManager.disabled = False
TaskManager.keep_running = None
TaskManager.task_list = []

def tearDown(self):
TaskManager.task_list = self.original_task_list

def test_enable_sets_disabled_false(self):
TaskManager.disabled = True
Expand Down Expand Up @@ -49,3 +55,100 @@ def test_good_stack_size_is_reasonable(self):

def test_start_new_thread_logs_warning(self):
TaskManager.start_new_thread()

def test_completed_task_is_removed_and_returns_result(self):
async def return_result():
return "result"

async def run_task():
task = TaskManager.create_task(return_result())
self.assertIn(task, TaskManager.task_list)
result = await task
return task, result

task, result = asyncio.run(run_task())

self.assertEqual(result, "result")
self.assertTrue(task.done())
self.assertTrue(task not in TaskManager.task_list)

def test_failed_task_is_removed_and_preserves_exception(self):
async def raise_error():
raise ValueError("failure")

async def run_task():
task = TaskManager.create_task(raise_error())
try:
await task
except ValueError as error:
return task, str(error)

task, message = asyncio.run(run_task())

self.assertEqual(message, "failure")
self.assertTrue(task.done())
self.assertTrue(task not in TaskManager.task_list)

def test_cancelled_task_is_removed_and_preserves_cancellation(self):
async def wait():
await asyncio.sleep(60)

async def run_task():
task = TaskManager.create_task(wait())
await asyncio.sleep_ms(0)
task.cancel()
try:
await task
except asyncio.CancelledError:
return task, True
return task, False

task, cancelled = asyncio.run(run_task())

self.assertTrue(cancelled)
self.assertTrue(task.done())
self.assertTrue(task not in TaskManager.task_list)

def test_cleanup_preserves_running_tasks(self):
async def wait_for_event(event):
await event.wait()

async def finish():
return None

async def run_tasks():
event = asyncio.Event()
running_task = TaskManager.create_task(wait_for_event(event))
completed_task = TaskManager.create_task(finish())
await completed_task
completed_removed = completed_task not in TaskManager.task_list
running_preserved = running_task in TaskManager.task_list
event.set()
await running_task
return completed_removed, running_preserved

completed_removed, running_preserved = asyncio.run(run_tasks())

self.assertTrue(completed_removed)
self.assertTrue(running_preserved)
self.assertEqual(TaskManager.task_list, [])

def test_many_completed_tasks_do_not_accumulate(self):
async def allocate_buffer():
buffer = bytearray(10 * 1024)
await asyncio.sleep_ms(0)
return len(buffer)

async def run_tasks():
tasks = []
for _ in range(100):
tasks.append(TaskManager.create_task(allocate_buffer()))
total = 0
for task in tasks:
total += await task
return total

total = asyncio.run(run_tasks())

self.assertEqual(total, 100 * 10 * 1024)
self.assertEqual(TaskManager.task_list, [])
Loading