thirdmagic 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. thirdmagic-0.0.1/.gitignore +231 -0
  2. thirdmagic-0.0.1/PKG-INFO +16 -0
  3. thirdmagic-0.0.1/pyproject.toml +53 -0
  4. thirdmagic-0.0.1/tests/__init__.py +0 -0
  5. thirdmagic-0.0.1/tests/unit/__init__.py +0 -0
  6. thirdmagic-0.0.1/tests/unit/assertions.py +80 -0
  7. thirdmagic-0.0.1/tests/unit/change_status/__init__.py +0 -0
  8. thirdmagic-0.0.1/tests/unit/change_status/assertions.py +14 -0
  9. thirdmagic-0.0.1/tests/unit/change_status/conftest.py +90 -0
  10. thirdmagic-0.0.1/tests/unit/change_status/test_chain.py +194 -0
  11. thirdmagic-0.0.1/tests/unit/change_status/test_signature.py +65 -0
  12. thirdmagic-0.0.1/tests/unit/change_status/test_swarm.py +230 -0
  13. thirdmagic-0.0.1/tests/unit/conftest.py +72 -0
  14. thirdmagic-0.0.1/tests/unit/creation/__init__.py +0 -0
  15. thirdmagic-0.0.1/tests/unit/creation/conftest.py +42 -0
  16. thirdmagic-0.0.1/tests/unit/creation/test_chain.py +219 -0
  17. thirdmagic-0.0.1/tests/unit/creation/test_chain_and_swarm.py +150 -0
  18. thirdmagic-0.0.1/tests/unit/creation/test_close_swarm.py +128 -0
  19. thirdmagic-0.0.1/tests/unit/creation/test_resolve_signature_keys.py +110 -0
  20. thirdmagic-0.0.1/tests/unit/creation/test_signature.py +269 -0
  21. thirdmagic-0.0.1/tests/unit/creation/test_swarm.py +60 -0
  22. thirdmagic-0.0.1/tests/unit/messages.py +24 -0
  23. thirdmagic-0.0.1/tests/unit/publish/__init__.py +0 -0
  24. thirdmagic-0.0.1/tests/unit/publish/conftest.py +38 -0
  25. thirdmagic-0.0.1/tests/unit/publish/test_aio_run_no_wait.py +65 -0
  26. thirdmagic-0.0.1/tests/unit/test_config.py +26 -0
  27. thirdmagic-0.0.1/tests/unit/utils.py +8 -0
  28. thirdmagic-0.0.1/thirdmagic/__init__.py +5 -0
  29. thirdmagic-0.0.1/thirdmagic/chain/__init__.py +4 -0
  30. thirdmagic-0.0.1/thirdmagic/chain/creator.py +34 -0
  31. thirdmagic-0.0.1/thirdmagic/chain/model.py +91 -0
  32. thirdmagic-0.0.1/thirdmagic/clients/__init__.py +3 -0
  33. thirdmagic-0.0.1/thirdmagic/clients/base.py +150 -0
  34. thirdmagic-0.0.1/thirdmagic/clients/lifecycle.py +24 -0
  35. thirdmagic-0.0.1/thirdmagic/consts.py +4 -0
  36. thirdmagic-0.0.1/thirdmagic/container.py +35 -0
  37. thirdmagic-0.0.1/thirdmagic/errors.py +26 -0
  38. thirdmagic-0.0.1/thirdmagic/message.py +26 -0
  39. thirdmagic-0.0.1/thirdmagic/signature/__init__.py +4 -0
  40. thirdmagic-0.0.1/thirdmagic/signature/model.py +191 -0
  41. thirdmagic-0.0.1/thirdmagic/signature/status.py +32 -0
  42. thirdmagic-0.0.1/thirdmagic/swarm/__init__.py +5 -0
  43. thirdmagic-0.0.1/thirdmagic/swarm/consts.py +1 -0
  44. thirdmagic-0.0.1/thirdmagic/swarm/creator.py +53 -0
  45. thirdmagic-0.0.1/thirdmagic/swarm/model.py +222 -0
  46. thirdmagic-0.0.1/thirdmagic/swarm/state.py +13 -0
  47. thirdmagic-0.0.1/thirdmagic/task/__init__.py +22 -0
  48. thirdmagic-0.0.1/thirdmagic/task/creator.py +87 -0
  49. thirdmagic-0.0.1/thirdmagic/task/model.py +83 -0
  50. thirdmagic-0.0.1/thirdmagic/task_def.py +12 -0
  51. thirdmagic-0.0.1/thirdmagic/typing_support.py +15 -0
  52. thirdmagic-0.0.1/thirdmagic/utils.py +57 -0
  53. thirdmagic-0.0.1/uv.lock +1352 -0
@@ -0,0 +1,231 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ #Pipfile.lock
96
+
97
+ # UV
98
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ #uv.lock
102
+
103
+ # poetry
104
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
105
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
106
+ # commonly ignored for libraries.
107
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
108
+ #poetry.lock
109
+ #poetry.toml
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
114
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
115
+ #pdm.lock
116
+ #pdm.toml
117
+ .pdm-python
118
+ .pdm-build/
119
+
120
+ # pixi
121
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
122
+ #pixi.lock
123
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
124
+ # in the .venv directory. It is recommended not to include this directory in version control.
125
+ .pixi
126
+
127
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
128
+ __pypackages__/
129
+
130
+ # Celery stuff
131
+ celerybeat-schedule
132
+ celerybeat.pid
133
+
134
+ # SageMath parsed files
135
+ *.sage.py
136
+
137
+ # Environments
138
+ .env
139
+ .envrc
140
+ .venv
141
+ env/
142
+ venv/
143
+ ENV/
144
+ env.bak/
145
+ venv.bak/
146
+
147
+ # Spyder project settings
148
+ .spyderproject
149
+ .spyproject
150
+
151
+ # Rope project settings
152
+ .ropeproject
153
+
154
+ # mkdocs documentation
155
+ /site
156
+ docs/roadmap.md
157
+
158
+ # mypy
159
+ .mypy_cache/
160
+ .dmypy.json
161
+ dmypy.json
162
+
163
+ # Pyre type checker
164
+ .pyre/
165
+
166
+ # pytype static type analyzer
167
+ .pytype/
168
+
169
+ # Cython debug symbols
170
+ cython_debug/
171
+
172
+ # PyCharm
173
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
174
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
175
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
176
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
177
+ #.idea/
178
+
179
+ # Abstra
180
+ # Abstra is an AI-powered process automation framework.
181
+ # Ignore directories containing user credentials, local state, and settings.
182
+ # Learn more at https://abstra.io/docs
183
+ .abstra/
184
+
185
+ # Visual Studio Code
186
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
187
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
188
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
189
+ # you could uncomment the following to ignore the entire vscode folder
190
+ # .vscode/
191
+
192
+ # Ruff stuff:
193
+ .ruff_cache/
194
+
195
+ # PyPI configuration file
196
+ .pypirc
197
+
198
+ # Cursor
199
+ # Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
200
+ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
201
+ # refer to https://docs.cursor.com/context/ignore-files
202
+ .cursorignore
203
+ .cursorindexingignore
204
+
205
+ # Marimo
206
+ marimo/_static/
207
+ marimo/_lsp/
208
+ __marimo__/
209
+ /.idea/
210
+ .idea/
211
+
212
+ /.junie/
213
+ CLAUDE.md
214
+ .vscode
215
+ future-features
216
+
217
+ # Ignore dynaconf secret files
218
+ .secrets.*
219
+ .claude
220
+ /junit.xml
221
+ .jbeval
222
+
223
+ # Mageflow visualizer static build (generated during CI)
224
+ mageflow/visualizer/static/*
225
+ !mageflow/visualizer/static/.gitkeep
226
+ app/node_modules/
227
+ frontend/node_modules/
228
+ /app/.idea/
229
+ /scripts
230
+ /frontend/.idea/
231
+ .planning/
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.4
2
+ Name: thirdmagic
3
+ Version: 0.0.1
4
+ Summary: Core models and signatures for mageflow task orchestration
5
+ Author-email: imaginary-cherry <yedidyakfir@gmail.com>
6
+ License: MIT
7
+ Requires-Python: <3.14,>=3.10
8
+ Requires-Dist: pydantic<3.0.0,>=2.0.0
9
+ Requires-Dist: rapyer<1.3.0,>=1.2.3
10
+ Provides-Extra: dev
11
+ Requires-Dist: black>=26.1.0; extra == 'dev'
12
+ Requires-Dist: coverage[toml]<8.0.0,>=7.0.0; extra == 'dev'
13
+ Requires-Dist: fakeredis[json,lua]<3.0.0,>=2.32.1; extra == 'dev'
14
+ Requires-Dist: hatchet-sdk>=1.23.0; extra == 'dev'
15
+ Requires-Dist: pytest-asyncio<2.0.0,>=1.2.0; extra == 'dev'
16
+ Requires-Dist: pytest<10.0.0,>=9.0.2; extra == 'dev'
@@ -0,0 +1,53 @@
1
+ [project]
2
+ name = "thirdmagic"
3
+ version = "0.0.1"
4
+ description = "Core models and signatures for mageflow task orchestration"
5
+ authors = [
6
+ {name = "imaginary-cherry", email = "yedidyakfir@gmail.com"}
7
+ ]
8
+ license = {text = "MIT"}
9
+ requires-python = ">=3.10,<3.14"
10
+ dependencies = [
11
+ "rapyer>=1.2.3,<1.3.0",
12
+ "pydantic>=2.0.0,<3.0.0",
13
+ ]
14
+ [project.optional-dependencies]
15
+ dev = [
16
+ "pytest>=9.0.2,<10.0.0",
17
+ "pytest-asyncio>=1.2.0,<2.0.0",
18
+ "fakeredis[json,lua]>=2.32.1,<3.0.0",
19
+ "coverage[toml]>=7.0.0,<8.0.0",
20
+ "hatchet-sdk>=1.23.0",
21
+ "black>=26.1.0",
22
+ ]
23
+
24
+ [build-system]
25
+ requires = ["hatchling"]
26
+ build-backend = "hatchling.build"
27
+
28
+ [tool.hatch.build.targets.wheel]
29
+ packages = ["thirdmagic"]
30
+
31
+ [tool.coverage.run]
32
+ source = ["thirdmagic"]
33
+ branch = true
34
+ omit = [
35
+ "*/tests/*",
36
+ "*/test_*"
37
+ ]
38
+
39
+ [tool.coverage.report]
40
+ skip_covered = false
41
+ show_missing = true
42
+ exclude_lines = [
43
+ "pragma: no cover",
44
+ "def __repr__",
45
+ "raise AssertionError",
46
+ "raise NotImplementedError",
47
+ "if __name__ == .__main__.:",
48
+ "@abstractmethod",
49
+ "@overload"
50
+ ]
51
+
52
+ [tool.coverage.html]
53
+ directory = "htmlcov"
File without changes
File without changes
@@ -0,0 +1,80 @@
1
+ from typing import TypeVar, Literal, cast
2
+
3
+ import rapyer
4
+ from rapyer.fields import RapyerKey
5
+ from redis.asyncio import Redis
6
+
7
+ from thirdmagic.consts import REMOVED_TASK_TTL
8
+ from thirdmagic.signature import Signature
9
+ from thirdmagic.task import TaskSignature
10
+
11
+ T = TypeVar("T", bound=Signature)
12
+ SwarmListName = Literal["finished_tasks", "failed_tasks", "tasks_results", "tasks"]
13
+
14
+
15
+ async def assert_task_reloaded_as_type(
16
+ task_key: RapyerKey,
17
+ expected_type: type[T],
18
+ ) -> T:
19
+ reloaded = await rapyer.aget(task_key)
20
+ assert reloaded is not None, f"Task {task_key} not found"
21
+ assert isinstance(
22
+ reloaded, expected_type
23
+ ), f"Expected {expected_type.__name__}, got {type(reloaded).__name__}"
24
+ return reloaded
25
+
26
+
27
+ def assert_callback_contains(
28
+ task: Signature,
29
+ success_keys: list[RapyerKey] | None = None,
30
+ error_keys: list[RapyerKey] | None = None,
31
+ ) -> None:
32
+ for success_key in success_keys or []:
33
+ assert (
34
+ success_key in task.success_callbacks
35
+ ), f"{success_key} not in success_callbacks"
36
+ for error_key in error_keys or []:
37
+ assert error_key in task.error_callbacks, f"{error_key} not in error_callbacks"
38
+
39
+
40
+ async def assert_tasks_not_exists(tasks_ids: list[str]):
41
+ for task_id in tasks_ids:
42
+ reloaded_signature = await TaskSignature.afind_one(task_id)
43
+ assert reloaded_signature is None
44
+
45
+
46
+ async def assert_tasks_changed_status(
47
+ tasks_ids: list[str | TaskSignature], status: str, old_status: str = None
48
+ ):
49
+ tasks_ids = tasks_ids if isinstance(tasks_ids, list) else [tasks_ids]
50
+ all_tasks = []
51
+ for task_key in tasks_ids:
52
+ task_key = task_key.key if isinstance(task_key, Signature) else task_key
53
+ reloaded_signature = await rapyer.aget(task_key)
54
+ reloaded_signature = cast(TaskSignature, reloaded_signature)
55
+ all_tasks.append(reloaded_signature)
56
+ assert reloaded_signature.task_status.status == status
57
+ if old_status:
58
+ assert reloaded_signature.task_status.last_status == old_status
59
+ return all_tasks
60
+
61
+
62
+ async def assert_redis_keys_do_not_contain_sub_task_ids(redis_client, sub_task_ids):
63
+ all_keys = await redis_client.keys("*")
64
+ all_keys_str = [
65
+ key.decode() if isinstance(key, bytes) else str(key) for key in all_keys
66
+ ]
67
+
68
+ for sub_task_id in sub_task_ids:
69
+ sub_task_id_str = str(sub_task_id)
70
+ keys_containing_sub_task = [
71
+ key for key in all_keys_str if sub_task_id_str in key
72
+ ]
73
+ assert (
74
+ not keys_containing_sub_task
75
+ ), f"Found Redis keys containing deleted sub-task ID {sub_task_id}: {keys_containing_sub_task}"
76
+
77
+
78
+ async def assert_task_has_short_ttl(redis_client: Redis, task_key: str):
79
+ ttl = await redis_client.ttl(task_key)
80
+ assert 0 < ttl <= REMOVED_TASK_TTL, f"Expected TTL <= {REMOVED_TASK_TTL}, got {ttl}"
File without changes
@@ -0,0 +1,14 @@
1
+ from typing import cast
2
+ from unittest.mock import Mock, call
3
+
4
+ from thirdmagic.task import TaskSignature
5
+
6
+
7
+ def assert_resume_signature(signature: TaskSignature, mock_adapter: Mock):
8
+ tasks_called = [a[0][0] for a in mock_adapter.acall_signature.call_args_list]
9
+ tasks_called = cast(list[TaskSignature], tasks_called)
10
+ task_called_with_id = [task for task in tasks_called if task.key == signature.key]
11
+ assert len(task_called_with_id) == 1, f"Task was resumed more than once"
12
+ mock_adapter.acall_signature.assert_has_awaits(
13
+ [call(task_called_with_id[0], None, set_return_field=False)]
14
+ )
@@ -0,0 +1,90 @@
1
+ from dataclasses import dataclass
2
+
3
+ import pytest
4
+ import pytest_asyncio
5
+
6
+ import thirdmagic
7
+ from tests.unit.messages import ContextMessage
8
+ from tests.unit.utils import extract_hatchet_validator
9
+ from thirdmagic.chain import ChainTaskSignature
10
+ from thirdmagic.swarm import SwarmTaskSignature
11
+ from thirdmagic.task import TaskSignature, SignatureStatus
12
+
13
+
14
+ @dataclass
15
+ class SwarmTestData:
16
+ task_signatures: list
17
+ swarm_signature: SwarmTaskSignature
18
+
19
+
20
+ @dataclass
21
+ class ChainTestData:
22
+ task_signatures: list
23
+ chain_signature: ChainTaskSignature
24
+
25
+
26
+ @dataclass
27
+ class TaskResumeConfig:
28
+ name: str
29
+ last_status: SignatureStatus
30
+
31
+
32
+ async def delete_tasks_by_indices(
33
+ task_signatures: list[TaskSignature],
34
+ indices: list[int],
35
+ ) -> list[str]:
36
+ deleted_task_ids = []
37
+ for idx in indices:
38
+ await task_signatures[idx].adelete()
39
+ deleted_task_ids.append(task_signatures[idx].key)
40
+ return deleted_task_ids
41
+
42
+
43
+ def get_non_deleted_task_keys(
44
+ task_signatures: list[TaskSignature],
45
+ deleted_indices: list[int],
46
+ ) -> list[str]:
47
+ return [
48
+ task_signatures[i].key
49
+ for i in range(len(task_signatures))
50
+ if i not in deleted_indices
51
+ ]
52
+
53
+
54
+ @pytest_asyncio.fixture
55
+ async def chain_with_tasks():
56
+ task_signatures = [
57
+ await thirdmagic.sign(f"chain_task_{i}", model_validators=ContextMessage)
58
+ for i in range(1, 4)
59
+ ]
60
+
61
+ chain_signature = await thirdmagic.chain([task.key for task in task_signatures])
62
+
63
+ return ChainTestData(
64
+ task_signatures=task_signatures, chain_signature=chain_signature
65
+ )
66
+
67
+
68
+ @pytest_asyncio.fixture
69
+ async def swarm_with_tasks():
70
+ task_signatures = [
71
+ await thirdmagic.sign(f"swarm_task_{i}", model_validators=ContextMessage)
72
+ for i in range(1, 4)
73
+ ]
74
+
75
+ swarm_signature = await thirdmagic.swarm(
76
+ task_name="test_swarm",
77
+ model_validators=ContextMessage,
78
+ tasks=task_signatures,
79
+ )
80
+
81
+ return SwarmTestData(
82
+ task_signatures=task_signatures, swarm_signature=swarm_signature
83
+ )
84
+
85
+
86
+ @pytest.fixture()
87
+ def hatchet_client_adapter(mock_adapter):
88
+ mock_adapter.extract_validator.side_effect = extract_hatchet_validator
89
+ mock_adapter.task_name.side_effect = lambda fn: fn.name
90
+ yield mock_adapter
@@ -0,0 +1,194 @@
1
+ import pytest
2
+
3
+ import thirdmagic
4
+ from tests.unit.assertions import (
5
+ assert_tasks_changed_status,
6
+ assert_tasks_not_exists,
7
+ assert_redis_keys_do_not_contain_sub_task_ids,
8
+ )
9
+ from tests.unit.change_status.assertions import assert_resume_signature
10
+ from tests.unit.change_status.conftest import (
11
+ TaskResumeConfig,
12
+ delete_tasks_by_indices,
13
+ get_non_deleted_task_keys,
14
+ )
15
+ from tests.unit.messages import ContextMessage
16
+ from thirdmagic.chain import ChainTaskSignature
17
+ from thirdmagic.task import TaskSignature, SignatureStatus
18
+
19
+
20
+ @pytest.mark.asyncio
21
+ async def test_chain_safe_change_status_on_deleted_signature_does_not_create_redis_entry_sanity():
22
+ # Arrange
23
+ task_signatures = [
24
+ await thirdmagic.sign(f"chain_task_{i}", model_validators=ContextMessage)
25
+ for i in range(1, 4)
26
+ ]
27
+ chain_signature = await thirdmagic.chain(
28
+ tasks=task_signatures, name="test_chain_unsaved"
29
+ )
30
+ chain_key = chain_signature.key
31
+ await chain_signature.adelete()
32
+
33
+ # Act
34
+ result = await ChainTaskSignature.safe_change_status(
35
+ chain_key, SignatureStatus.SUSPENDED
36
+ )
37
+
38
+ # Assert
39
+ assert result is False
40
+ reloaded_signature = await TaskSignature.afind_one(chain_key)
41
+ assert reloaded_signature is None
42
+
43
+
44
+ @pytest.mark.asyncio
45
+ @pytest.mark.parametrize(
46
+ ["task_names", "tasks_to_delete_indices", "new_status"],
47
+ [
48
+ [
49
+ ["task1", "task2", "task3"],
50
+ [],
51
+ SignatureStatus.SUSPENDED,
52
+ ],
53
+ [
54
+ ["task1", "task2"],
55
+ [0, 1],
56
+ SignatureStatus.CANCELED,
57
+ ],
58
+ [
59
+ ["task1", "task2", "task3"],
60
+ [0, 2],
61
+ SignatureStatus.ACTIVE,
62
+ ],
63
+ ],
64
+ )
65
+ async def test_chain_change_status_with_optional_deleted_sub_tasks_edge_case(
66
+ redis_client,
67
+ task_names: list[str],
68
+ tasks_to_delete_indices: list[int],
69
+ new_status: SignatureStatus,
70
+ mock_task_def,
71
+ ):
72
+ # Arrange
73
+ task_signatures = [await thirdmagic.sign(name) for name in task_names]
74
+ chain_signature = await thirdmagic.chain([task.key for task in task_signatures])
75
+ deleted_task_ids = await delete_tasks_by_indices(
76
+ task_signatures, tasks_to_delete_indices
77
+ )
78
+
79
+ # Act
80
+ await chain_signature.safe_change_status(chain_signature.key, new_status)
81
+
82
+ # Assert
83
+ reloaded_chain = await TaskSignature.aget(chain_signature.key)
84
+ assert reloaded_chain.task_status.status == new_status
85
+ assert reloaded_chain.task_status.last_status == SignatureStatus.PENDING
86
+
87
+ await assert_tasks_not_exists(deleted_task_ids)
88
+
89
+ non_deleted_keys = get_non_deleted_task_keys(
90
+ task_signatures, tasks_to_delete_indices
91
+ )
92
+ await assert_tasks_changed_status(
93
+ non_deleted_keys, new_status, SignatureStatus.PENDING
94
+ )
95
+
96
+ await assert_redis_keys_do_not_contain_sub_task_ids(redis_client, deleted_task_ids)
97
+
98
+
99
+ @pytest.mark.asyncio
100
+ @pytest.mark.parametrize(
101
+ ["task_configs", "tasks_to_delete_indices"],
102
+ [
103
+ [
104
+ [
105
+ TaskResumeConfig(name="task1", last_status=SignatureStatus.ACTIVE),
106
+ TaskResumeConfig(name="task2", last_status=SignatureStatus.ACTIVE),
107
+ TaskResumeConfig(name="task3", last_status=SignatureStatus.ACTIVE),
108
+ ],
109
+ [],
110
+ ],
111
+ [
112
+ [
113
+ TaskResumeConfig(name="task1", last_status=SignatureStatus.PENDING),
114
+ TaskResumeConfig(name="task2", last_status=SignatureStatus.PENDING),
115
+ ],
116
+ [0],
117
+ ],
118
+ [
119
+ [
120
+ TaskResumeConfig(name="task1", last_status=SignatureStatus.ACTIVE),
121
+ TaskResumeConfig(name="task2", last_status=SignatureStatus.ACTIVE),
122
+ TaskResumeConfig(name="task3", last_status=SignatureStatus.ACTIVE),
123
+ ],
124
+ [1, 2],
125
+ ],
126
+ ],
127
+ )
128
+ async def test_chain_resume_with_optional_deleted_sub_tasks_sanity(
129
+ mock_adapter,
130
+ task_configs: list[TaskResumeConfig],
131
+ tasks_to_delete_indices: list[int],
132
+ mock_task_def,
133
+ ):
134
+ # Arrange
135
+ task_signatures = []
136
+ expected_statuses = []
137
+ num_of_aio_run = 0
138
+ for config in task_configs:
139
+ task_signature = await thirdmagic.sign(config.name)
140
+ task_signature.task_status.status = SignatureStatus.SUSPENDED
141
+ task_signature.task_status.last_status = config.last_status
142
+ await task_signature.asave()
143
+ task_signatures.append(task_signature)
144
+ expected_statuses.append(config.last_status)
145
+
146
+ chain_signature = await thirdmagic.chain([task.key for task in task_signatures])
147
+ chain_signature.task_status.status = SignatureStatus.SUSPENDED
148
+
149
+ deleted_task_ids = await delete_tasks_by_indices(
150
+ task_signatures, tasks_to_delete_indices
151
+ )
152
+
153
+ # Act
154
+ await chain_signature.resume()
155
+
156
+ # Assert
157
+ non_deleted_task_indices = [
158
+ i for i in range(len(task_signatures)) if i not in tasks_to_delete_indices
159
+ ]
160
+ for i in non_deleted_task_indices:
161
+ task = task_signatures[i]
162
+ new_status = expected_statuses[i]
163
+ if new_status == SignatureStatus.ACTIVE:
164
+ new_status = SignatureStatus.PENDING
165
+ assert_resume_signature(task, mock_adapter)
166
+ num_of_aio_run += 1
167
+ await assert_tasks_changed_status(
168
+ [task.key], new_status, SignatureStatus.SUSPENDED
169
+ )
170
+
171
+ await assert_tasks_changed_status(
172
+ [chain_signature.key], SignatureStatus.PENDING, SignatureStatus.SUSPENDED
173
+ )
174
+
175
+ await assert_tasks_not_exists(deleted_task_ids)
176
+ assert mock_adapter.acall_signature.call_count == num_of_aio_run
177
+
178
+
179
+ @pytest.mark.asyncio
180
+ async def test_chain_suspend_sanity(chain_with_tasks):
181
+ # Arrange
182
+ chain_data = chain_with_tasks
183
+
184
+ # Act
185
+ await chain_data.chain_signature.suspend()
186
+
187
+ # Assert
188
+ # Verify all tasks changed status to suspend
189
+ await assert_tasks_changed_status(
190
+ [chain_data.chain_signature.key], SignatureStatus.SUSPENDED
191
+ )
192
+ await assert_tasks_changed_status(
193
+ [task.key for task in chain_data.task_signatures], SignatureStatus.SUSPENDED
194
+ )