mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-09-21 02:16:27 +00:00
b0fd0a9fa3
* feat(cli): track checkpoint command and TUI usage * feat(telemetry): track runtime checkpoint operations * fix(telemetry): count prune usage after argument validation * style(cli): format checkpoint prune command --------- Co-authored-by: Vidit Ostwal <110953813+Vidit-Ostwal@users.noreply.github.com>
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
"""Runtime checkpoint operations report feature usage, including Python-only calls."""
|
|
|
|
from collections import Counter
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from crewai.events.event_bus import crewai_event_bus
|
|
from crewai.events.event_listener import event_listener
|
|
from crewai.events.types.checkpoint_events import (
|
|
CheckpointCompletedEvent,
|
|
CheckpointFailedEvent,
|
|
CheckpointForkCompletedEvent,
|
|
CheckpointPrunedEvent,
|
|
CheckpointRestoreCompletedEvent,
|
|
CheckpointRestoreFailedEvent,
|
|
)
|
|
from crewai.state.checkpoint_config import CheckpointConfig
|
|
from crewai.state.checkpoint_listener import _do_checkpoint
|
|
from crewai.state.provider.json_provider import JsonProvider
|
|
from crewai.state.runtime import RuntimeState
|
|
|
|
from ..utils import wait_for_event_handlers
|
|
|
|
|
|
@pytest.fixture
|
|
def features(monkeypatch):
|
|
recorded = []
|
|
with crewai_event_bus.scoped_handlers():
|
|
event_listener.setup_listeners(crewai_event_bus)
|
|
monkeypatch.setattr(
|
|
event_listener._telemetry, "feature_usage_span", recorded.append
|
|
)
|
|
yield recorded
|
|
wait_for_event_handlers()
|
|
|
|
|
|
def test_manual_save_restore_and_fork(tmp_path, features):
|
|
state = RuntimeState(root=[])
|
|
location = state.checkpoint(str(tmp_path))
|
|
restored = RuntimeState.from_checkpoint(CheckpointConfig(restore_from=location))
|
|
restored.fork("private-branch")
|
|
wait_for_event_handlers()
|
|
|
|
assert Counter(features) == {
|
|
"checkpoint:save": 1,
|
|
"checkpoint:restore": 1,
|
|
"checkpoint:fork": 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_save_and_restore(tmp_path, features):
|
|
state = RuntimeState(root=[])
|
|
location = await state.acheckpoint(str(tmp_path))
|
|
await RuntimeState.afrom_checkpoint(CheckpointConfig(restore_from=location))
|
|
wait_for_event_handlers()
|
|
|
|
assert Counter(features) == {"checkpoint:save": 1, "checkpoint:restore": 1}
|
|
|
|
|
|
def test_automatic_checkpoint_and_prune(tmp_path, features):
|
|
_do_checkpoint(
|
|
RuntimeState(root=[]),
|
|
CheckpointConfig(location=str(tmp_path), max_checkpoints=1),
|
|
)
|
|
wait_for_event_handlers()
|
|
|
|
assert Counter(features) == {"checkpoint:save": 1, "checkpoint:prune": 1}
|
|
|
|
|
|
def test_failed_operations_do_not_count_as_success(tmp_path, features):
|
|
with patch.object(JsonProvider, "checkpoint", side_effect=OSError("private path")):
|
|
with pytest.raises(OSError):
|
|
RuntimeState(root=[]).checkpoint(str(tmp_path))
|
|
with pytest.raises(FileNotFoundError):
|
|
RuntimeState.from_checkpoint(
|
|
CheckpointConfig(restore_from=str(tmp_path / "missing.json"))
|
|
)
|
|
wait_for_event_handlers()
|
|
|
|
assert Counter(features) == {
|
|
"checkpoint:save_failed": 1,
|
|
"checkpoint:restore_failed": 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"event",
|
|
[
|
|
CheckpointCompletedEvent(
|
|
location="private",
|
|
provider="JsonProvider",
|
|
checkpoint_id="private-id",
|
|
duration_ms=1,
|
|
),
|
|
CheckpointFailedEvent(
|
|
location="private", provider="JsonProvider", error="private error"
|
|
),
|
|
CheckpointRestoreCompletedEvent(
|
|
location="private", checkpoint_id="private-id", duration_ms=1
|
|
),
|
|
CheckpointRestoreFailedEvent(location="private", error="private error"),
|
|
CheckpointForkCompletedEvent(branch="private-branch"),
|
|
CheckpointPrunedEvent(
|
|
location="private",
|
|
provider="JsonProvider",
|
|
removed_count=1,
|
|
max_checkpoints=1,
|
|
),
|
|
],
|
|
)
|
|
def test_replayed_checkpoint_events_are_not_counted(event, features):
|
|
future = crewai_event_bus.replay(None, event)
|
|
if future is not None:
|
|
future.result(timeout=5)
|
|
wait_for_event_handlers()
|
|
|
|
assert features == []
|