Coverage for utils_aiosqlite/tests/unit/test_task_manager_str.py: 26%
179 statements
« prev ^ index » next coverage.py v7.10.7, created at 2025-10-30 14:44 +0800
1"""Unit tests for TaskManager with string task IDs."""
3import tempfile
4from collections.abc import AsyncGenerator
5from pathlib import Path
6from unittest.mock import patch
8import orjson
9import pandas as pd
10import pytest
12from utils_aiosqlite import TaskManager
15class TestTaskManagerStr:
16 """Test TaskManager with string task IDs."""
18 @pytest.fixture
19 async def task_manager(self) -> AsyncGenerator[TaskManager[str], None]:
20 """Create a TaskManager instance with str task IDs for testing."""
21 with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
22 db_path = Path(tmp.name)
24 manager = await TaskManager.async_init(db_path, column_type="str")
25 yield manager
27 # Cleanup
28 if db_path.exists():
29 db_path.unlink()
31 @pytest.mark.asyncio
32 async def test_async_init_str_default(self):
33 """Test TaskManager initialization with default str column type."""
34 with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
35 db_path = Path(tmp.name)
37 try:
38 # Default should be "str"
39 manager = await TaskManager.async_init(db_path)
40 assert manager.pth_tasks_db == db_path
41 assert manager.column_type == "str"
42 assert manager.default_timeout == 5
43 finally:
44 if db_path.exists():
45 db_path.unlink()
47 @pytest.mark.asyncio
48 async def test_async_init_str_explicit(self):
49 """Test TaskManager initialization with explicit str column type."""
50 with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
51 db_path = Path(tmp.name)
53 try:
54 manager = await TaskManager.async_init(db_path, column_type="str")
55 assert manager.pth_tasks_db == db_path
56 assert manager.column_type == "str"
57 assert manager.default_timeout == 5
58 finally:
59 if db_path.exists():
60 db_path.unlink()
62 @pytest.mark.asyncio
63 async def test_setup_tasks_str(self, task_manager: TaskManager[str]):
64 """Test setting up tasks with string IDs."""
65 task_ids = [("task1",), ("task2",), ("task3",)]
66 await task_manager.setup_tasks(task_ids)
68 all_tasks = await task_manager.get_all_tasks()
69 assert sorted(all_tasks) == ["task1", "task2", "task3"]
71 @pytest.mark.asyncio
72 async def test_setup_tasks_with_special_chars(self, task_manager: TaskManager[str]):
73 """Test setting up tasks with string IDs containing special characters."""
74 task_ids = [("task-1",), ("task_2",), ("task.3",), ("task@4",)]
75 await task_manager.setup_tasks(task_ids)
77 all_tasks = await task_manager.get_all_tasks()
78 assert sorted(all_tasks) == ["task-1", "task.3", "task@4", "task_2"]
80 @pytest.mark.asyncio
81 async def test_update_outputs_str(self, task_manager: TaskManager[str]):
82 """Test updating outputs with string task IDs."""
83 await task_manager.setup_tasks([("task1",), ("task2",)])
85 outputs = [{"result": "success", "value": 42}, {"result": "error", "message": "failed"}]
86 await task_manager.update_outputs(["task1", "task2"], outputs)
88 output_1 = await task_manager.get_output("task1")
89 output_2 = await task_manager.get_output("task2")
91 assert output_1 == {"result": "success", "value": 42}
92 assert output_2 == {"result": "error", "message": "failed"}
94 @pytest.mark.asyncio
95 async def test_update_outputs_n_metadatas_str(self, task_manager: TaskManager[str]):
96 """Test updating outputs and metadata with string task IDs."""
97 await task_manager.setup_tasks([("task1",), ("task2",)])
99 outputs = [{"result": "success"}, {"result": "error"}]
100 metadatas = [{"duration": 1.5}, {"error_code": 500}]
102 await task_manager.update_outputs_n_metadatas(["task1", "task2"], outputs, metadatas)
104 output_1 = await task_manager.get_output("task1")
105 metadata_1 = await task_manager.get_metadata("task1")
106 output_2 = await task_manager.get_output("task2")
107 metadata_2 = await task_manager.get_metadata("task2")
109 assert output_1 == {"result": "success"}
110 assert metadata_1 == {"duration": 1.5}
111 assert output_2 == {"result": "error"}
112 assert metadata_2 == {"error_code": 500}
114 @pytest.mark.asyncio
115 async def test_get_output_nonexistent_str(self, task_manager: TaskManager[str]):
116 """Test getting output for non-existent string task ID."""
117 result = await task_manager.get_output("nonexistent")
118 assert result is None
120 @pytest.mark.asyncio
121 async def test_get_metadata_nonexistent_str(self, task_manager: TaskManager[str]):
122 """Test getting metadata for non-existent string task ID."""
123 result = await task_manager.get_metadata("nonexistent")
124 assert result is None
126 @pytest.mark.asyncio
127 async def test_get_all_output_str(self, task_manager: TaskManager[str]):
128 """Test getting all outputs with string task IDs."""
129 await task_manager.setup_tasks([("task1",), ("task2",)])
130 outputs = [{"result": "success"}, {"result": "error"}]
131 await task_manager.update_outputs(["task1", "task2"], outputs)
133 df = await task_manager.get_all_output()
134 assert len(df) == 2
135 assert set(df["task_id"]) == {"task1", "task2"}
136 assert df[df["task_id"] == "task1"]["output"].iloc[0] == {"result": "success"}
137 assert df[df["task_id"] == "task2"]["output"].iloc[0] == {"result": "error"}
139 @pytest.mark.asyncio
140 async def test_get_all_output_flatten_str(self, task_manager: TaskManager[str]):
141 """Test getting all outputs with flattening enabled."""
142 await task_manager.setup_tasks([("task1",), ("task2",)])
143 outputs = [{"result": "success", "count": 5}, {"result": "error", "count": 0}]
144 await task_manager.update_outputs(["task1", "task2"], outputs)
146 df = await task_manager.get_all_output(flatten=True)
147 assert len(df) == 2
148 assert "result" in df.columns
149 assert "count" in df.columns
150 assert set(df["task_id"]) == {"task1", "task2"}
152 @pytest.mark.asyncio
153 async def test_get_all_metadata_str(self, task_manager: TaskManager[str]):
154 """Test getting all metadata with string task IDs."""
155 await task_manager.setup_tasks([("task1",), ("task2",)])
156 outputs = [{"result": "success"}, {"result": "error"}]
157 metadatas = [{"duration": 1.5}, {"error_code": 500}]
158 await task_manager.update_outputs_n_metadatas(["task1", "task2"], outputs, metadatas)
160 df = await task_manager.get_all_metadata()
161 assert len(df) == 2
162 assert set(df["task_id"]) == {"task1", "task2"}
163 assert df[df["task_id"] == "task1"]["metadata"].iloc[0] == {"duration": 1.5}
164 assert df[df["task_id"] == "task2"]["metadata"].iloc[0] == {"error_code": 500}
166 @pytest.mark.asyncio
167 async def test_get_pending_tasks_str(self, task_manager: TaskManager[str]):
168 """Test getting pending tasks with string IDs."""
169 await task_manager.setup_tasks([("task1",), ("task2",), ("task3",)])
170 await task_manager.update_outputs(["task1"], [{"result": "success"}])
172 pending = await task_manager.get_pending_tasks()
173 assert sorted(pending) == ["task2", "task3"]
175 @pytest.mark.asyncio
176 async def test_get_noresult_tasks_str(self, task_manager: TaskManager[str]):
177 """Test getting no-result tasks with string IDs."""
178 await task_manager.setup_tasks([("task1",), ("task2",), ("task3",)])
179 await task_manager.update_outputs(["task1", "task2"], [{"result": "success"}, {}])
181 noresult = await task_manager.get_noresult_tasks()
182 assert noresult == ["task2"]
184 @pytest.mark.asyncio
185 async def test_mark_noresult_tasks_as_pending_str(self, task_manager: TaskManager[str]):
186 """Test marking no-result tasks as pending."""
187 await task_manager.setup_tasks([("task1",), ("task2",)])
188 await task_manager.update_outputs_n_metadatas(
189 ["task1", "task2"], [{"result": "success"}, {}], [{"meta": "data"}, {"meta": "data2"}]
190 )
192 # Verify task2 has empty output but metadata
193 output_2 = await task_manager.get_output("task2")
194 metadata_2 = await task_manager.get_metadata("task2")
195 assert output_2 == {}
196 assert metadata_2 == {"meta": "data2"}
198 await task_manager.mark_noresult_tasks_as_pending()
200 # After marking as pending, both output and metadata should be None
201 output_2_after = await task_manager.get_output("task2")
202 metadata_2_after = await task_manager.get_metadata("task2")
203 assert output_2_after is None
204 assert metadata_2_after is None
206 pending = await task_manager.get_pending_tasks()
207 assert "task2" in pending
209 @pytest.mark.asyncio
210 async def test_query_str(self, task_manager: TaskManager[str]):
211 """Test custom query execution with string task IDs."""
212 await task_manager.setup_tasks([("task1",), ("task2",), ("task3",)])
214 df = await task_manager.query("SELECT COUNT(*) as count FROM tasks")
215 assert len(df) == 1
216 assert df["count"].iloc[0] == 3
218 # Test query with WHERE clause on string task_id
219 df = await task_manager.query("SELECT task_id FROM tasks WHERE task_id LIKE 'task%'")
220 assert len(df) == 3
221 assert set(df["task_id"]) == {"task1", "task2", "task3"}
223 @pytest.mark.asyncio
224 async def test_executemany_str(self, task_manager: TaskManager[str]):
225 """Test executemany method with string task IDs."""
226 # Create some initial tasks
227 await task_manager.setup_tasks([("task1",), ("task2",), ("task3",)])
229 # Use executemany to update multiple tasks at once
230 data = [(orjson.dumps({"result": f"batch_{task_id}"}), task_id) for task_id in ["task1", "task2", "task3"]]
231 await task_manager.executemany("UPDATE tasks SET output = ? WHERE task_id = ?", data)
233 # Verify updates
234 for task_id in ["task1", "task2", "task3"]:
235 output = await task_manager.get_output(task_id)
236 assert output == {"result": f"batch_{task_id}"}
238 @pytest.mark.asyncio
239 async def test_mixed_string_formats(self, task_manager: TaskManager[str]):
240 """Test TaskManager with various string formats."""
241 # Test with different string formats
242 task_ids = [
243 ("uuid-123-456",),
244 ("task_with_underscores",),
245 ("task-with-dashes",),
246 ("task.with.dots",),
247 ("UPPERCASE_TASK",),
248 ("MixedCase_Task",),
249 ("123_numeric_prefix",),
250 ("task_with_numbers_123",),
251 ]
253 await task_manager.setup_tasks(task_ids)
255 all_tasks = await task_manager.get_all_tasks()
256 expected_tasks = [tid[0] for tid in task_ids]
257 assert sorted(all_tasks) == sorted(expected_tasks)
259 # Test updating some of these tasks
260 test_outputs = [{"status": "completed", "type": "uuid"}, {"status": "failed", "type": "underscore"}]
262 await task_manager.update_outputs(["uuid-123-456", "task_with_underscores"], test_outputs)
264 # Verify the updates
265 output_uuid = await task_manager.get_output("uuid-123-456")
266 output_underscore = await task_manager.get_output("task_with_underscores")
268 assert output_uuid == {"status": "completed", "type": "uuid"}
269 assert output_underscore == {"status": "failed", "type": "underscore"}
271 @pytest.mark.asyncio
272 async def test_long_string_task_ids(self, task_manager: TaskManager[str]):
273 """Test TaskManager with very long string task IDs."""
274 long_task_id = "very_long_task_id_" + "x" * 200 # 218 characters total
276 await task_manager.setup_tasks([(long_task_id,)])
278 # Test that we can work with long task IDs
279 await task_manager.update_outputs([long_task_id], [{"result": "success_with_long_id"}])
281 output = await task_manager.get_output(long_task_id)
282 assert output == {"result": "success_with_long_id"}
284 all_tasks = await task_manager.get_all_tasks()
285 assert long_task_id in all_tasks