From 3c1719f001d5f54be969a989b5646b7109fd1b98 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 11:23:38 +0200 Subject: [PATCH 1/2] Make Workspace name lower cased --- src/codegate/db/models.py | 22 +++++++++++----------- src/codegate/pipeline/cli/commands.py | 4 ++-- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/codegate/db/models.py b/src/codegate/db/models.py index 6120ea1f..ac54eacc 100644 --- a/src/codegate/db/models.py +++ b/src/codegate/db/models.py @@ -1,8 +1,7 @@ import datetime -import re -from typing import Any, Optional +from typing import Annotated, Any, Optional -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, StringConstraints class Alert(BaseModel): @@ -40,18 +39,19 @@ class Setting(BaseModel): other_settings: Optional[Any] +WorskpaceNameStr = Annotated[ + str, + StringConstraints( + strip_whitespace=True, to_lower=True, pattern=r"^[a-zA-Z0-9_-]+$", strict=True + ), +] + + class Workspace(BaseModel): id: str - name: str + name: WorskpaceNameStr system_prompt: Optional[str] - @field_validator("name", mode="plain") - @classmethod - def name_must_be_alphanumeric(cls, value): - if not re.match(r"^[a-zA-Z0-9_-]+$", value): - raise ValueError("name must be alphanumeric and can only contain _ and -") - return value - class Session(BaseModel): id: str diff --git a/src/codegate/pipeline/cli/commands.py b/src/codegate/pipeline/cli/commands.py index 597aa5e5..24797be8 100644 --- a/src/codegate/pipeline/cli/commands.py +++ b/src/codegate/pipeline/cli/commands.py @@ -185,7 +185,7 @@ async def _add_workspace(self, flags: Dict[str, str], args: List[str]) -> str: return "Please provide a name. Use `codegate workspace add your_workspace_name`" try: - _ = await self.workspace_crud.add_workspace(new_workspace_name) + ws = await self.workspace_crud.add_workspace(new_workspace_name) except ValidationError: return "Invalid workspace name: It should be alphanumeric and dashes" except AlreadyExistsError: @@ -195,7 +195,7 @@ async def _add_workspace(self, flags: Dict[str, str], args: List[str]) -> str: except Exception: return "An error occurred while adding the workspace" - return f"Workspace **{new_workspace_name}** has been added" + return f"Workspace **{ws.name}** has been added" async def _rename_workspace(self, flags: Dict[str, str], args: List[str]) -> str: """ From cffe8e03e4ed25a679a2f1718622c5cc4aa68377 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Wed, 22 Jan 2025 11:39:03 +0200 Subject: [PATCH 2/2] Fix unit tests --- tests/pipeline/workspace/test_workspace.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/pipeline/workspace/test_workspace.py b/tests/pipeline/workspace/test_workspace.py index 27db0519..891ecf8a 100644 --- a/tests/pipeline/workspace/test_workspace.py +++ b/tests/pipeline/workspace/test_workspace.py @@ -2,6 +2,7 @@ import pytest +from codegate.db.models import Workspace as WorkspaceModel from codegate.db.models import WorkspaceActive from codegate.pipeline.cli.commands import Workspace @@ -80,7 +81,8 @@ async def test_add_workspaces(args, existing_workspaces, expected_message): with patch("codegate.workspaces.crud.WorkspaceCrud", autospec=True) as mock_recorder_cls: mock_recorder = mock_recorder_cls.return_value workspace_commands.workspace_crud = mock_recorder - mock_recorder.add_workspace = AsyncMock() + created_workspace = WorkspaceModel(id="1", name="myworkspace", system_prompt=None) + mock_recorder.add_workspace = AsyncMock(return_value=created_workspace) # Call the method result = await workspace_commands._add_workspace(None, args)