Skip to content
14 changes: 14 additions & 0 deletions autogpt_platform/backend/backend/blocks/time_blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -445,6 +445,8 @@ class Input(BlockSchemaInput):
repeat: int = SchemaField(
description="Number of times to repeat the timer",
default=1,
ge=1,
le=1000,
Comment thread
kcze marked this conversation as resolved.
)

class Output(BlockSchemaOutput):
Expand All @@ -469,6 +471,8 @@ def __init__(self):
],
)

MAX_TOTAL_SECONDS = 7 * 86400 # 7 days
Comment thread
kcze marked this conversation as resolved.
Outdated

async def run(self, input_data: Input, **kwargs) -> BlockOutput:
seconds = int(input_data.seconds)
minutes = int(input_data.minutes)
Expand All @@ -477,6 +481,16 @@ async def run(self, input_data: Input, **kwargs) -> BlockOutput:

total_seconds = seconds + minutes * 60 + hours * 3600 + days * 86400

if total_seconds < 0:
raise ValueError(
f"Countdown duration must be non-negative, got {total_seconds}s"
)
if total_seconds > self.MAX_TOTAL_SECONDS:
raise ValueError(
f"Countdown duration {total_seconds}s exceeds max "
f"({self.MAX_TOTAL_SECONDS}s = 7 days)"
)
Comment thread
kcze marked this conversation as resolved.
Outdated

for _ in range(input_data.repeat):
if total_seconds > 0:
await asyncio.sleep(total_seconds)
Expand Down
37 changes: 37 additions & 0 deletions autogpt_platform/backend/backend/blocks/time_blocks_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
from unittest.mock import AsyncMock

import pytest

from backend.blocks.time_blocks import CountdownTimerBlock


async def _run(block: CountdownTimerBlock, **input_kwargs):
outputs = []
async for name, value in block.run(block.input_schema(**input_kwargs)):
outputs.append((name, value))
return outputs


@pytest.mark.asyncio
async def test_countdown_timer_rejects_excessive_duration():
block = CountdownTimerBlock()
with pytest.raises(ValueError, match="exceeds max"):
await _run(block, days=365, repeat=1)


@pytest.mark.asyncio
async def test_countdown_timer_rejects_negative_duration():
block = CountdownTimerBlock()
with pytest.raises(ValueError, match="non-negative"):
await _run(block, seconds="-1")


@pytest.mark.asyncio
async def test_countdown_timer_allows_duration_at_cap(mocker):
sleep_mock = mocker.patch(
"backend.blocks.time_blocks.asyncio.sleep", new_callable=AsyncMock
)
block = CountdownTimerBlock()
outputs = await _run(block, days=7, repeat=1)
assert outputs == [("output_message", "timer finished")]
sleep_mock.assert_awaited_once_with(7 * 86400)
Loading