Skip to content

Commit 55a4754

Browse files
NiveditJainNiveditJain
committed
Add unit tests for prune and re-enqueue signal functionality
- Introduced comprehensive tests for the new prune and re-enqueue signal routes, validating request models and response handling. - Implemented model validation tests for PruneRequestModel and ReEnqueueAfterRequestModel, ensuring correct data handling. - Added controller tests for prune_signal and re_queue_after_signal functions, covering success and error scenarios. - Enhanced SignalResponseModel tests to verify correct serialization and deserialization of responses. - Improved overall test coverage for state management operations, ensuring robustness and reliability.
1 parent c689a8e commit 55a4754

4 files changed

Lines changed: 1160 additions & 3 deletions

File tree

‎state-manager/tests/unit/controller/test_prune_signal.py‎

Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
import pytest
2+
from unittest.mock import AsyncMock, MagicMock, patch
3+
from fastapi import HTTPException, status
4+
from beanie import PydanticObjectId
5+
6+
from app.controller.prune_signal import prune_signal
7+
from app.models.signal_models import PruneRequestModel, SignalResponseModel
8+
from app.models.state_status_enum import StateStatusEnum
9+
10+
11+
class TestPruneSignal:
12+
"""Test cases for prune_signal function"""
13+
14+
@pytest.fixture
15+
def mock_request_id(self):
16+
return "test-request-id"
17+
18+
@pytest.fixture
19+
def mock_namespace(self):
20+
return "test_namespace"
21+
22+
@pytest.fixture
23+
def mock_state_id(self):
24+
return PydanticObjectId()
25+
26+
@pytest.fixture
27+
def mock_prune_request(self):
28+
return PruneRequestModel(
29+
data={"key": "value", "nested": {"data": "test"}}
30+
)
31+
32+
@pytest.fixture
33+
def mock_state_created(self):
34+
state = MagicMock()
35+
state.id = PydanticObjectId()
36+
state.status = StateStatusEnum.CREATED
37+
state.enqueue_after = 1234567890
38+
return state
39+
40+
@patch('app.controller.prune_signal.State')
41+
async def test_prune_signal_success(
42+
self,
43+
mock_state_class,
44+
mock_namespace,
45+
mock_state_id,
46+
mock_prune_request,
47+
mock_state_created,
48+
mock_request_id
49+
):
50+
"""Test successful pruning of state"""
51+
# Arrange
52+
mock_state_created.save = AsyncMock()
53+
mock_state_class.find_one = AsyncMock(return_value=mock_state_created)
54+
55+
# Act
56+
result = await prune_signal(
57+
mock_namespace,
58+
mock_state_id,
59+
mock_prune_request,
60+
mock_request_id
61+
)
62+
63+
# Assert
64+
assert result.status == StateStatusEnum.PRUNED
65+
assert result.enqueue_after == 1234567890
66+
assert mock_state_created.status == StateStatusEnum.PRUNED
67+
assert mock_state_created.data == mock_prune_request.data
68+
assert mock_state_created.save.call_count == 1
69+
assert mock_state_class.find_one.call_count == 1
70+
71+
@patch('app.controller.prune_signal.State')
72+
async def test_prune_signal_state_not_found(
73+
self,
74+
mock_state_class,
75+
mock_namespace,
76+
mock_state_id,
77+
mock_prune_request,
78+
mock_request_id
79+
):
80+
"""Test when state is not found"""
81+
# Arrange
82+
mock_state_class.find_one = AsyncMock(return_value=None)
83+
84+
# Act & Assert
85+
with pytest.raises(HTTPException) as exc_info:
86+
await prune_signal(
87+
mock_namespace,
88+
mock_state_id,
89+
mock_prune_request,
90+
mock_request_id
91+
)
92+
93+
assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND
94+
assert exc_info.value.detail == "State not found"
95+
96+
@patch('app.controller.prune_signal.State')
97+
async def test_prune_signal_invalid_status_queued(
98+
self,
99+
mock_state_class,
100+
mock_namespace,
101+
mock_state_id,
102+
mock_prune_request,
103+
mock_request_id
104+
):
105+
"""Test when state is in QUEUED status (invalid for pruning)"""
106+
# Arrange
107+
mock_state = MagicMock()
108+
mock_state.status = StateStatusEnum.QUEUED
109+
mock_state_class.find_one = AsyncMock(return_value=mock_state)
110+
111+
# Act & Assert
112+
with pytest.raises(HTTPException) as exc_info:
113+
await prune_signal(
114+
mock_namespace,
115+
mock_state_id,
116+
mock_prune_request,
117+
mock_request_id
118+
)
119+
120+
assert exc_info.value.status_code == status.HTTP_400_BAD_REQUEST
121+
assert exc_info.value.detail == "State is not created"
122+
123+
@patch('app.controller.prune_signal.State')
124+
async def test_prune_signal_invalid_status_executed(
125+
self,
126+
mock_state_class,
127+
mock_namespace,
128+
mock_state_id,
129+
mock_prune_request,
130+
mock_request_id
131+
):
132+
"""Test when state is in EXECUTED status (invalid for pruning)"""
133+
# Arrange
134+
mock_state = MagicMock()
135+
mock_state.status = StateStatusEnum.EXECUTED
136+
mock_state_class.find_one = AsyncMock(return_value=mock_state)
137+
138+
# Act & Assert
139+
with pytest.raises(HTTPException) as exc_info:
140+
await prune_signal(
141+
mock_namespace,
142+
mock_state_id,
143+
mock_prune_request,
144+
mock_request_id
145+
)
146+
147+
assert exc_info.value.status_code == status.HTTP_400_BAD_REQUEST
148+
assert exc_info.value.detail == "State is not created"
149+
150+
@patch('app.controller.prune_signal.State')
151+
async def test_prune_signal_invalid_status_errored(
152+
self,
153+
mock_state_class,
154+
mock_namespace,
155+
mock_state_id,
156+
mock_prune_request,
157+
mock_request_id
158+
):
159+
"""Test when state is in ERRORED status (invalid for pruning)"""
160+
# Arrange
161+
mock_state = MagicMock()
162+
mock_state.status = StateStatusEnum.ERRORED
163+
mock_state_class.find_one = AsyncMock(return_value=mock_state)
164+
165+
# Act & Assert
166+
with pytest.raises(HTTPException) as exc_info:
167+
await prune_signal(
168+
mock_namespace,
169+
mock_state_id,
170+
mock_prune_request,
171+
mock_request_id
172+
)
173+
174+
assert exc_info.value.status_code == status.HTTP_400_BAD_REQUEST
175+
assert exc_info.value.detail == "State is not created"
176+
177+
@patch('app.controller.prune_signal.State')
178+
async def test_prune_signal_invalid_status_pruned(
179+
self,
180+
mock_state_class,
181+
mock_namespace,
182+
mock_state_id,
183+
mock_prune_request,
184+
mock_request_id
185+
):
186+
"""Test when state is already in PRUNED status (invalid for pruning)"""
187+
# Arrange
188+
mock_state = MagicMock()
189+
mock_state.status = StateStatusEnum.PRUNED
190+
mock_state_class.find_one = AsyncMock(return_value=mock_state)
191+
192+
# Act & Assert
193+
with pytest.raises(HTTPException) as exc_info:
194+
await prune_signal(
195+
mock_namespace,
196+
mock_state_id,
197+
mock_prune_request,
198+
mock_request_id
199+
)
200+
201+
assert exc_info.value.status_code == status.HTTP_400_BAD_REQUEST
202+
assert exc_info.value.detail == "State is not created"
203+
204+
@patch('app.controller.prune_signal.State')
205+
async def test_prune_signal_database_error(
206+
self,
207+
mock_state_class,
208+
mock_namespace,
209+
mock_state_id,
210+
mock_prune_request,
211+
mock_request_id
212+
):
213+
"""Test handling of database errors"""
214+
# Arrange
215+
mock_state_class.find_one = MagicMock(side_effect=Exception("Database error"))
216+
217+
# Act & Assert
218+
with pytest.raises(Exception) as exc_info:
219+
await prune_signal(
220+
mock_namespace,
221+
mock_state_id,
222+
mock_prune_request,
223+
mock_request_id
224+
)
225+
226+
assert str(exc_info.value) == "Database error"
227+
228+
@patch('app.controller.prune_signal.State')
229+
async def test_prune_signal_save_error(
230+
self,
231+
mock_state_class,
232+
mock_namespace,
233+
mock_state_id,
234+
mock_prune_request,
235+
mock_state_created,
236+
mock_request_id
237+
):
238+
"""Test handling of save errors"""
239+
# Arrange
240+
mock_state_created.save = AsyncMock(side_effect=Exception("Save error"))
241+
mock_state_class.find_one = AsyncMock(return_value=mock_state_created)
242+
243+
# Act & Assert
244+
with pytest.raises(Exception) as exc_info:
245+
await prune_signal(
246+
mock_namespace,
247+
mock_state_id,
248+
mock_prune_request,
249+
mock_request_id
250+
)
251+
252+
assert str(exc_info.value) == "Save error"
253+
254+
@patch('app.controller.prune_signal.State')
255+
async def test_prune_signal_with_empty_data(
256+
self,
257+
mock_state_class,
258+
mock_namespace,
259+
mock_state_id,
260+
mock_state_created,
261+
mock_request_id
262+
):
263+
"""Test pruning with empty data"""
264+
# Arrange
265+
prune_request = PruneRequestModel(data={})
266+
mock_state_created.save = AsyncMock()
267+
mock_state_class.find_one = AsyncMock(return_value=mock_state_created)
268+
269+
# Act
270+
result = await prune_signal(
271+
mock_namespace,
272+
mock_state_id,
273+
prune_request,
274+
mock_request_id
275+
)
276+
277+
# Assert
278+
assert result.status == StateStatusEnum.PRUNED
279+
assert mock_state_created.data == {}
280+
assert mock_state_created.save.call_count == 1
281+
282+
@patch('app.controller.prune_signal.State')
283+
async def test_prune_signal_with_complex_data(
284+
self,
285+
mock_state_class,
286+
mock_namespace,
287+
mock_state_id,
288+
mock_state_created,
289+
mock_request_id
290+
):
291+
"""Test pruning with complex nested data"""
292+
# Arrange
293+
complex_data = {
294+
"string": "test",
295+
"number": 42,
296+
"boolean": True,
297+
"list": [1, 2, 3],
298+
"nested": {
299+
"object": {
300+
"deep": "value"
301+
}
302+
}
303+
}
304+
prune_request = PruneRequestModel(data=complex_data)
305+
mock_state_created.save = AsyncMock()
306+
mock_state_class.find_one = AsyncMock(return_value=mock_state_created)
307+
308+
# Act
309+
result = await prune_signal(
310+
mock_namespace,
311+
mock_state_id,
312+
prune_request,
313+
mock_request_id
314+
)
315+
316+
# Assert
317+
assert result.status == StateStatusEnum.PRUNED
318+
assert mock_state_created.data == complex_data
319+
assert mock_state_created.save.call_count == 1

0 commit comments

Comments
 (0)