Files

726 lines
21 KiB
Python

"""Step definitions for domain models testing."""
from datetime import datetime
from decimal import Decimal
from pathlib import Path
from behave import given, then, when
from pydantic import ValidationError
from cleveragents.domain.models.aimodels_custom import CustomAIModel
from cleveragents.domain.models.aimodelsdatamodels import AIModel, Provider
from cleveragents.domain.models.core import (
Change,
ChangeSet,
CloudBillingFields,
Context,
ContextType,
ContextUpdateResult,
CreditsTransaction,
CreditsTransactionType,
CreditType,
Invite,
OperationType,
Org,
OrgRole,
OrgUser,
Plan,
PlanStatus,
Project,
ProjectSettings,
SummaryForUpdateContextParams,
User,
)
from cleveragents.domain.models.core.enums import ModelProvider
from cleveragents.domain.models.orguserconfig import OrgUserConfig
from cleveragents.domain.models.planconfig import PlanConfig
def _sample_plan_config() -> PlanConfig:
"""Create a reusable PlanConfig instance for tests."""
return PlanConfig(
auto_mode="manual",
editor="vim",
editor_command="vim",
editor_args=[],
editor_open_manually=False,
auto_continue=False,
auto_build=False,
auto_update_context=False,
auto_context=False,
smart_context=False,
auto_apply=False,
auto_commit=False,
skip_commit=False,
can_exec=False,
auto_exec=False,
auto_debug=False,
auto_debug_tries=0,
auto_revert_on_rewind=False,
skip_changes_menu=False,
)
@given("I have valid project data")
def step_valid_project_data(context):
"""Set up valid project data."""
context.project_data = {
"name": "test-project",
"path": Path("/tmp/test-project"),
"settings": ProjectSettings(),
}
@given('I have project data with invalid name "{name}"')
def step_invalid_project_name(context, name):
"""Set up project data with invalid name."""
context.project_data = {
"name": name,
"path": Path("/tmp/test-project"),
}
@when("I create a Project model")
def step_create_project(context):
"""Create a Project model from data."""
try:
context.project = Project(**context.project_data)
context.error = None
except ValidationError as e:
context.project = None
context.error = e
@when("I try to create a Project model")
def step_try_create_project(context):
"""Try to create a Project model, expecting it might fail."""
try:
context.project = Project(**context.project_data)
context.error = None
except ValidationError as e:
context.project = None
context.error = e
@then("the project model should be created successfully")
def step_project_created(context):
"""Check that project was created."""
assert context.project is not None
assert context.error is None
@then("the project name should be validated")
def step_project_name_validated(context):
"""Check that project name is valid."""
assert context.project.name == context.project_data["name"]
@then("the project path should be absolute")
def step_project_path_absolute(context):
"""Check that project path is absolute."""
assert context.project.path.is_absolute()
@then("a validation error should be raised")
def step_validation_error_raised(context):
"""Check that validation error was raised."""
assert context.error is not None
assert isinstance(context.error, ValidationError)
@then('the error should mention "{message}"')
def step_error_message(context, message):
"""Check that error contains expected message."""
assert message in str(context.error)
# Plan-related steps
@given("I have valid plan data")
def step_valid_plan_data(context):
"""Set up valid plan data."""
context.plan_data = {
"project_id": 1,
"name": "test-plan",
"prompt": "Add error handling",
"status": PlanStatus.PENDING,
}
@when("I create a Plan model")
def step_create_plan(context):
"""Create a Plan model from data."""
context.plan = Plan(**context.plan_data)
@then("the plan model should be created successfully")
def step_plan_created(context):
"""Check that plan was created."""
assert context.plan is not None
@then('the plan status should be "{status}"')
def step_plan_status(context, status):
"""Check plan status."""
assert context.plan.status == status
@then("the plan should have timestamps")
def step_plan_timestamps(context):
"""Check that plan has timestamps."""
assert isinstance(context.plan.created_at, datetime)
assert isinstance(context.plan.updated_at, datetime)
@given('I have a plan with status "{status}"')
def step_plan_with_status(context, status):
"""Create a plan with specific status."""
context.plan = Plan(
project_id=1,
name="test-plan",
prompt="Test",
status=PlanStatus(status),
)
@when('I update the plan status to "{status}"')
def step_update_plan_status(context, status):
"""Update plan status."""
context.plan.status = PlanStatus(status)
@then('the updated plan status should be "{status}"')
def step_check_updated_plan_status(context, status):
"""Check updated plan status."""
assert context.plan.status == status
@then("the updated_at timestamp should be updated")
def step_check_updated_timestamp(context):
"""Check that updated_at is set."""
# In real implementation, this would be handled by the repository
assert context.plan.updated_at is not None
# Context-related steps
@given("I have valid context data")
def step_valid_context_data(context):
"""Set up valid context data."""
context.context_data = {
"plan_id": 1,
"type": ContextType.FILE,
"path": "/tmp/test.py",
"size": 1024,
}
@when("I create a Context model")
def step_create_context(context):
"""Create a Context model from data."""
context.context_obj = Context(**context.context_data)
@then("the context model should be created successfully")
def step_context_created(context):
"""Check that context was created."""
assert context.context_obj is not None
@then('the context type should be "{type_name}"')
def step_context_type(context, type_name):
"""Check context type."""
assert context.context_obj.type == type_name
@then("the context should have a path")
def step_context_path(context):
"""Check that context has a path."""
assert context.context_obj.path is not None
assert len(context.context_obj.path) > 0
@given("I have context update statistics")
def step_context_update_stats(context):
"""Prepare representative context update statistics."""
context.context_update_data = {
"updatedContexts": [
{
"plan_id": 1,
"type": ContextType.FILE,
"path": "/tmp/main.py",
"size": 512,
},
{
"plan_id": 1,
"type": ContextType.URL,
"path": "https://example.com/docs",
"size": 0,
},
],
"tokenDiffsById": {
"ctx-1": {"added": 120, "removed": 20},
"ctx-2": {"added": 0, "removed": 0},
},
"tokensDiff": 100,
"totalTokens": 640,
"numFiles": 1,
"numUrls": 1,
"numImages": 0,
"numTrees": 0,
"numMaps": 0,
"maxTokens": 4096,
}
@when("I create a ContextUpdateResult summary")
def step_create_context_update_result(context):
"""Instantiate a ContextUpdateResult model from prepared data."""
context.context_update_result = ContextUpdateResult(**context.context_update_data)
@then("the context update summary should report correct totals")
def step_context_update_totals(context):
"""Verify the aggregated counts returned by ContextUpdateResult."""
result = context.context_update_result
assert len(result.updated_contexts) == 2
assert result.num_files == 1
assert result.num_urls == 1
assert result.num_images == 0
assert result.tokens_diff == 100
assert result.total_tokens == 640
@then("the token diffs should remain non-negative")
def step_context_update_token_diff_non_negative(context):
"""Ensure nested token diff counters never drop below zero."""
result = context.context_update_result
for diff_map in result.token_diffs_by_id.values():
assert all(value >= 0 for value in diff_map.values())
@given("I have summary counts for context updates")
def step_summary_counts_context_updates(context):
"""Prepare summary counts mimicking camelCase payloads."""
context.context_summary_data = {
"numFiles": 3,
"numTrees": 1,
"numUrls": 2,
"numMaps": 0,
"tokensDiff": 25,
"totalTokens": 512,
}
@when("I create a SummaryForUpdateContextParams model")
def step_create_summary_update_context_params(context):
"""Create SummaryForUpdateContextParams from provided counts."""
context.context_summary = SummaryForUpdateContextParams(
**context.context_summary_data
)
@then("the summary counts should match the source data")
def step_verify_summary_update_context(context):
"""Validate that camelCase keys hydrate correctly into the model."""
summary = context.context_summary
assert summary.num_files == 3
assert summary.num_trees == 1
assert summary.num_urls == 2
assert summary.num_maps == 0
assert summary.tokens_diff == 25
assert summary.total_tokens == 512
# Change-related steps
@given("I have valid change data")
def step_valid_change_data(context):
"""Set up valid change data."""
context.change_data = {
"plan_id": 1,
"file_path": "/tmp/test.py",
"operation": OperationType.MODIFY,
"original_content": "old content",
"new_content": "new content",
}
@when("I create a Change model")
def step_create_change(context):
"""Create a Change model from data."""
context.change = Change(**context.change_data)
@then("the change model should be created successfully")
def step_change_created(context):
"""Check that change was created."""
assert context.change is not None
@then("the change operation should be valid")
def step_change_operation_valid(context):
"""Check that change operation is valid."""
assert context.change.operation in OperationType
@then("the change should not be applied by default")
def step_change_not_applied(context):
"""Check that change is not applied by default."""
assert context.change.applied is False
# ChangeSet-related steps
@given("I have a changeset with multiple changes")
def step_changeset_with_changes(context):
"""Create a changeset with multiple changes."""
changes = [
Change(
plan_id=1,
file_path="/tmp/file1.py",
operation=OperationType.CREATE,
new_content="new file",
),
Change(
plan_id=1,
file_path="/tmp/file2.py",
operation=OperationType.MODIFY,
original_content="old",
new_content="new",
),
Change(
plan_id=1,
file_path="/tmp/file3.py",
operation=OperationType.DELETE,
original_content="deleted content",
),
]
context.changeset = ChangeSet(plan_id=1, changes=changes)
@when("I get the changeset statistics")
def step_get_changeset_stats(context):
"""Get statistics from changeset."""
context.stats = context.changeset.stats
@then("the stats should show correct counts")
def step_stats_correct_counts(context):
"""Check that stats show correct counts."""
assert context.stats["total"] == 3
assert context.stats["creates"] == 1
assert context.stats["modifies"] == 1
assert context.stats["deletes"] == 1
@then("the stats should count creates, modifies, and deletes")
def step_stats_count_operations(context):
"""Check that stats count all operation types."""
assert "creates" in context.stats
assert "modifies" in context.stats
assert "deletes" in context.stats
assert "applied" in context.stats
# Domain models from cloud/organization data
def _sample_billing_payload() -> dict[str, object]:
"""Provide representative camelCase billing payload."""
return {
"creditsBalance": "125.5",
"monthlyGrant": "200.0",
"autoRebuyEnabled": True,
"autoRebuyMinThreshold": "5.0",
"autoRebuyToBalance": "50.0",
"notifyThreshold": "10.0",
"maxThresholdPerMonth": "500.0",
"billingCycleStartedAt": "2025-01-01T00:00:00Z",
"changedBillingMode": False,
"trialPaid": True,
"stripeSubscriptionId": "sub-123",
"subscriptionStatus": "active",
"subscriptionPausedAt": "2025-01-10T00:00:00Z",
"stripePaymentMethod": "pm_abc",
"subscriptionActionRequired": False,
"subscriptionActionRequiredInvoiceUrl": None,
}
@given("I have sample cloud billing data")
def step_sample_cloud_billing(context):
"""Prepare representative cloud billing data."""
context.cloud_billing_data = _sample_billing_payload()
@when("I create a CloudBillingFields model")
def step_create_billing_fields(context):
"""Instantiate CloudBillingFields from camelCase payload."""
context.billing_fields = CloudBillingFields(**context.cloud_billing_data)
@then("the billing fields should expose decimal values")
def step_verify_billing_fields(context):
"""Validate decimal conversion for billing fields."""
billing = context.billing_fields
assert billing.credits_balance == Decimal("125.5")
assert billing.monthly_grant == Decimal("200.0")
assert billing.subscription_action_required is False
@given("I have organization data with billing")
def step_org_with_billing(context):
"""Prepare organization payload referencing billing fields."""
context.org_payload = {
"id": "org-1",
"name": "Example Org",
"isTrial": True,
"autoAddDomainUsers": True,
"integratedModelsMode": False,
}
context.org_billing_payload = _sample_billing_payload()
@when("I create an Org domain model")
def step_create_org_model(context):
"""Instantiate Org with nested billing fields."""
billing = CloudBillingFields(**context.org_billing_payload)
context.org_model = Org(cloud_billing_fields=billing, **context.org_payload)
@then("the org should include billing configuration")
def step_verify_org_model(context):
"""Ensure the org retains billing configuration."""
org = context.org_model
assert org.cloud_billing_fields is not None
assert org.cloud_billing_fields.notify_threshold == Decimal("10.0")
@given("I have an org role definition")
def step_org_role_definition(context):
"""Set up org role payload."""
context.org_role_data = {
"id": "role-1",
"isDefault": True,
"label": "Owner",
"description": "Full access",
}
@when("I create an OrgRole model")
def step_create_org_role(context):
"""Instantiate OrgRole."""
context.org_role = OrgRole(**context.org_role_data)
@then("the org role should retain its label")
def step_verify_org_role(context):
"""Ensure org role values persist."""
assert context.org_role.label == "Owner"
assert context.org_role.is_default is True
@given("I have user data referencing a default plan config")
def step_user_with_plan_config(context):
"""Prepare user payload referencing plan config."""
context.user_plan_config = _sample_plan_config()
context.user_data = {
"id": "user-1",
"name": "Casey",
"email": "casey@example.com",
"isTrial": False,
"numNonDraftPlans": 3,
}
@when("I create a User model")
def step_create_user_model(context):
"""Instantiate User with nested plan config."""
context.user_model = User(
default_plan_config=context.user_plan_config, **context.user_data
)
@then("the user should include the default plan config")
def step_verify_user_model(context):
"""Ensure user retains plan config."""
assert context.user_model.default_plan_config is not None
assert context.user_model.default_plan_config.auto_mode == "manual"
@given("I have an org user association")
def step_org_user_association(context):
"""Prepare OrgUser payload."""
context.org_user_config = OrgUserConfig(
prompted_claude_max=False,
use_claude_subscription=True,
claude_subscription_cooldown_started_at=datetime(2025, 1, 1, 0, 0, 0),
)
context.org_user_data = {
"orgId": "org-1",
"userId": "user-1",
"orgRoleId": "role-1",
}
@when("I create an OrgUser model")
def step_create_org_user(context):
"""Instantiate OrgUser."""
context.org_user_model = OrgUser(
config=context.org_user_config, **context.org_user_data
)
@then("the org user config should be attached")
def step_verify_org_user(context):
"""Ensure OrgUser carries config payload."""
assert context.org_user_model.config is not None
assert context.org_user_model.config.use_claude_subscription is True
@given("I have invite metadata")
def step_invite_metadata(context):
"""Prepare Invite payload."""
context.invite_data = {
"id": "invite-1",
"orgId": "org-1",
"email": "new@example.com",
"name": "New User",
"orgRoleId": "role-1",
"inviterId": "user-1",
"inviteeId": None,
"acceptedAt": "2025-02-01T12:00:00Z",
"createdAt": "2025-01-15T08:30:00Z",
}
@when("I create an Invite model")
def step_create_invite(context):
"""Instantiate Invite."""
context.invite_model = Invite(**context.invite_data)
@then("the invite should track acceptance timestamps")
def step_verify_invite(context):
"""Ensure Invite datetimes parsed."""
assert isinstance(context.invite_model.accepted_at, datetime)
assert context.invite_model.accepted_at.year == 2025
@given("I have credits transaction data")
def step_credits_transaction_data(context):
"""Prepare CreditsTransaction payload."""
context.credits_transaction_data = {
"id": "txn-1",
"orgId": "org-1",
"orgName": "Example Org",
"userId": "user-1",
"userEmail": "casey@example.com",
"userName": "Casey",
"transactionType": CreditsTransactionType.CREDIT,
"amount": "42.0",
"startBalance": "20.0",
"endBalance": "62.0",
"creditType": CreditType.TRIAL,
"creditIsAutoRebuy": True,
"creditAutoRebuyMinThreshold": "5.0",
"creditAutoRebuyToBalance": "40.0",
"debitModelProvider": ModelProvider.OPENAI,
"debitModelRole": "ModelRolePlanner",
"debitModelName": "gpt-4.1",
"debitPlanId": "plan-1",
"debitPlanName": "Core Plan",
"createdAt": "2025-01-20T00:00:00Z",
}
@when("I create a CreditsTransaction model")
def step_create_credits_transaction(context):
"""Instantiate CreditsTransaction."""
context.credits_transaction = CreditsTransaction(**context.credits_transaction_data)
@then("the credits transaction should track balances")
def step_verify_credits_transaction(context):
"""Ensure CreditsTransaction stores decimals and provider metadata."""
tx = context.credits_transaction
assert tx.amount == Decimal("42.0")
assert tx.debit_model_provider == ModelProvider.OPENAI
assert tx.end_balance == Decimal("62.0")
# Auto-generated models steps
@when("I import the auto-generated models")
def step_import_autogenerated(context):
"""Import auto-generated models."""
try:
from cleveragents.domain.models.auth import AuthHeader
from cleveragents.domain.models.planconfig import PlanConfig
from cleveragents.domain.models.stream import BuildInfo
context.import_error = None
context.imported_models = [AuthHeader, BuildInfo, PlanConfig]
except Exception as e:
context.import_error = e
context.imported_models = []
@then("all models should load without errors")
def step_models_load_without_errors(context):
"""Check that models loaded without errors."""
assert context.import_error is None
assert len(context.imported_models) > 0
@then("the models should have proper Pydantic configuration")
def step_models_have_pydantic_config(context):
"""Check that models have Pydantic configuration."""
for model_class in context.imported_models:
assert hasattr(model_class, "model_config")
assert hasattr(model_class, "model_validate")
@given('I create an "{model_name}" with the following data:')
@given('I create a "{model_name}" with the following data:')
def create_model(context, model_name):
data = {row["name"]: row["value"] for row in context.table}
if model_name == "AIModel":
context.model = AIModel(**data)
elif model_name == "Provider":
context.model = Provider(**data)
elif model_name == "CustomAIModel":
context.model = CustomAIModel(**data)
@then('the "{model_name}" object should have the following attributes:')
def verify_model_attributes(context, model_name):
for row in context.table:
attr_name = row["name"]
expected_value = row["value"]
actual_value = getattr(context.model, attr_name)
if expected_value == "False":
expected_value = False
assert actual_value == expected_value, (
f"Expected {attr_name} to be {expected_value}, but got {actual_value}"
)