Coverage for utils_aiosqlite/tests/unit/test_task_manager_int.py: 27%
164 statements
« prev ^ index » next coverage.py v7.10.7, created at 2025-10-30 14:44 +0800
1"""Unit tests for TaskManager with integer 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 TestTaskManagerInt:
16 """Test TaskManager with integer task IDs."""
18 @pytest.fixture
19 async def task_manager(self) -> AsyncGenerator[TaskManager[int], None]:
20 """Create a TaskManager instance with int 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="int")
25 yield manager
27 # Cleanup
28 if db_path.exists():
29 db_path.unlink()
31 @pytest.mark.asyncio
32 async def test_async_init_int(self):
33 """Test TaskManager initialization with int column type."""
34 with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
35 db_path = Path(tmp.name)
37 try:
38 manager = await TaskManager.async_init(db_path, column_type="int")
39 assert manager.pth_tasks_db == db_path
40 assert manager.column_type == "int"
41 assert manager.default_timeout == 5
42 finally:
43 if db_path.exists():
44 db_path.unlink()
46 @pytest.mark.asyncio
47 async def test_async_init_invalid_column_type(self):
48 """Test TaskManager initialization with invalid column type."""
49 with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
50 db_path = Path(tmp.name)
52 try:
53 with pytest.raises(ValueError, match="column_type must be either 'int' or 'str'"):
54 await TaskManager.async_init(db_path, column_type="invalid")
55 finally:
56 if db_path.exists():
57 db_path.unlink()
59 @pytest.mark.asyncio
60 async def test_setup_tasks_int(self, task_manager: TaskManager[int]):
61 """Test setting up tasks with integer IDs."""
62 task_ids = [(1,), (2,), (3,)]
63 await task_manager.setup_tasks(task_ids)
65 all_tasks = await task_manager.get_all_tasks()
66 assert sorted(all_tasks) == [1, 2, 3]
68 @pytest.mark.asyncio
69 async def test_setup_tasks_empty_list(self, task_manager: TaskManager[int]):
70 """Test setup_tasks with empty list raises ValueError."""
71 with pytest.raises(ValueError, match="task_id_tpls cannot be empty"):
72 await task_manager.setup_tasks([])
74 @pytest.mark.asyncio
75 async def test_update_outputs_int(self, task_manager: TaskManager[int]):
76 """Test updating outputs with integer task IDs."""
77 await task_manager.setup_tasks([(1,), (2,)])
79 outputs = [{"result": "success", "value": 42}, {"result": "error", "message": "failed"}]
80 await task_manager.update_outputs([1, 2], outputs)
82 output_1 = await task_manager.get_output(1)
83 output_2 = await task_manager.get_output(2)
85 assert output_1 == {"result": "success", "value": 42}
86 assert output_2 == {"result": "error", "message": "failed"}
88 @pytest.mark.asyncio
89 async def test_update_outputs_n_metadatas_int(self, task_manager: TaskManager[int]):
90 """Test updating outputs and metadata with integer task IDs."""
91 await task_manager.setup_tasks([(1,), (2,)])
93 outputs = [{"result": "success"}, {"result": "error"}]
94 metadatas = [{"duration": 1.5}, {"error_code": 500}]
96 await task_manager.update_outputs_n_metadatas([1, 2], outputs, metadatas)
98 output_1 = await task_manager.get_output(1)
99 metadata_1 = await task_manager.get_metadata(1)
100 output_2 = await task_manager.get_output(2)
101 metadata_2 = await task_manager.get_metadata(2)
103 assert output_1 == {"result": "success"}
104 assert metadata_1 == {"duration": 1.5}
105 assert output_2 == {"result": "error"}
106 assert metadata_2 == {"error_code": 500}
108 @pytest.mark.asyncio
109 async def test_get_output_nonexistent_int(self, task_manager: TaskManager[int]):
110 """Test getting output for non-existent integer task ID."""
111 result = await task_manager.get_output(999)
112 assert result is None
114 @pytest.mark.asyncio
115 async def test_get_metadata_nonexistent_int(self, task_manager: TaskManager[int]):
116 """Test getting metadata for non-existent integer task ID."""
117 result = await task_manager.get_metadata(999)
118 assert result is None
120 @pytest.mark.asyncio
121 async def test_get_all_output_int(self, task_manager: TaskManager[int]):
122 """Test getting all outputs with integer task IDs."""
123 await task_manager.setup_tasks([(1,), (2,)])
124 outputs = [{"result": "success"}, {"result": "error"}]
125 await task_manager.update_outputs([1, 2], outputs)
127 df = await task_manager.get_all_output()
128 assert len(df) == 2
129 assert set(df["task_id"]) == {1, 2}
130 assert df[df["task_id"] == 1]["output"].iloc[0] == {"result": "success"}
131 assert df[df["task_id"] == 2]["output"].iloc[0] == {"result": "error"}
133 @pytest.mark.asyncio
134 async def test_get_all_output_flatten_int(self, task_manager: TaskManager[int]):
135 """Test getting all outputs with flattening enabled."""
136 await task_manager.setup_tasks([(1,), (2,)])
137 outputs = [{"result": "success", "count": 5}, {"result": "error", "count": 0}]
138 await task_manager.update_outputs([1, 2], outputs)
140 df = await task_manager.get_all_output(flatten=True)
141 assert len(df) == 2
142 assert "result" in df.columns
143 assert "count" in df.columns
144 assert set(df["task_id"]) == {1, 2}
146 @pytest.mark.asyncio
147 async def test_get_all_metadata_int(self, task_manager: TaskManager[int]):
148 """Test getting all metadata with integer task IDs."""
149 await task_manager.setup_tasks([(1,), (2,)])
150 outputs = [{"result": "success"}, {"result": "error"}]
151 metadatas = [{"duration": 1.5}, {"error_code": 500}]
152 await task_manager.update_outputs_n_metadatas([1, 2], outputs, metadatas)
154 df = await task_manager.get_all_metadata()
155 assert len(df) == 2
156 assert set(df["task_id"]) == {1, 2}
157 assert df[df["task_id"] == 1]["metadata"].iloc[0] == {"duration": 1.5}
158 assert df[df["task_id"] == 2]["metadata"].iloc[0] == {"error_code": 500}
160 @pytest.mark.asyncio
161 async def test_get_pending_tasks_int(self, task_manager: TaskManager[int]):
162 """Test getting pending tasks with integer IDs."""
163 await task_manager.setup_tasks([(1,), (2,), (3,)])
164 await task_manager.update_outputs([1], [{"result": "success"}])
166 pending = await task_manager.get_pending_tasks()
167 assert sorted(pending) == [2, 3]
169 @pytest.mark.asyncio
170 async def test_get_noresult_tasks_int(self, task_manager: TaskManager[int]):
171 """Test getting no-result tasks with integer IDs."""
172 await task_manager.setup_tasks([(1,), (2,), (3,)])
173 await task_manager.update_outputs([1, 2], [{"result": "success"}, {}])
175 noresult = await task_manager.get_noresult_tasks()
176 assert noresult == [2]
178 @pytest.mark.asyncio
179 async def test_mark_noresult_tasks_as_pending_int(self, task_manager: TaskManager[int]):
180 """Test marking no-result tasks as pending."""
181 await task_manager.setup_tasks([(1,), (2,)])
182 await task_manager.update_outputs_n_metadatas(
183 [1, 2], [{"result": "success"}, {}], [{"meta": "data"}, {"meta": "data2"}]
184 )
186 # Verify task 2 has empty output but metadata
187 output_2 = await task_manager.get_output(2)
188 metadata_2 = await task_manager.get_metadata(2)
189 assert output_2 == {}
190 assert metadata_2 == {"meta": "data2"}
192 await task_manager.mark_noresult_tasks_as_pending()
194 # After marking as pending, both output and metadata should be None
195 output_2_after = await task_manager.get_output(2)
196 metadata_2_after = await task_manager.get_metadata(2)
197 assert output_2_after is None
198 assert metadata_2_after is None
200 pending = await task_manager.get_pending_tasks()
201 assert 2 in pending
203 @pytest.mark.asyncio
204 async def test_query_int(self, task_manager: TaskManager[int]):
205 """Test custom query execution."""
206 await task_manager.setup_tasks([(1,), (2,), (3,)])
208 df = await task_manager.query("SELECT COUNT(*) as count FROM tasks")
209 assert len(df) == 1
210 assert df["count"].iloc[0] == 3
212 @pytest.mark.asyncio
213 async def test_static_methods(self):
214 """Test static JSON serialization methods."""
215 test_dict = {"key": "value", "number": 42}
217 # Test _dumps
218 blob = TaskManager._dumps(test_dict)
219 assert isinstance(blob, bytes)
221 # Test _loads
222 result = TaskManager._loads(blob)
223 assert result == test_dict
225 # Test _loads with invalid data
226 invalid_result = TaskManager._loads(b"invalid json")
227 assert invalid_result is None
229 # Test _dumps with unserializable data (should return None)
230 class UnserializableClass:
231 pass
233 unserializable = {"obj": UnserializableClass()}
234 result = TaskManager._dumps(unserializable)
235 assert result is None
237 @pytest.mark.asyncio
238 async def test_executemany_int(self, task_manager: TaskManager[int]):
239 """Test executemany method."""
240 # Create some initial tasks
241 await task_manager.setup_tasks([(1,), (2,), (3,)])
243 # Use executemany to update multiple tasks at once
244 data = [(orjson.dumps({"result": f"batch_{i}"}), i) for i in [1, 2, 3]]
245 await task_manager.executemany("UPDATE tasks SET output = ? WHERE task_id = ?", data)
247 # Verify updates
248 for i in [1, 2, 3]:
249 output = await task_manager.get_output(i)
250 assert output == {"result": f"batch_{i}"}