Fixing issues on quickstart and adding tests

Author: aniket-s-kulkarni

.github/workflows/pytest.yml

@@ -0,0 +1,29 @@
+name: Run Pytest with uv
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Install uv
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.13'
+
+      - name: Run tests
+        run: |
+          uv run pytest tests# if not in requirements.txt

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

README.md

@@ -49,7 +49,7 @@ architecture-beta
 
 The MCP server runs locally on the machine that runs the LLM frontend (eg Claude). The installation steps are simple
 
-1. Clone or download this repository. 
+1. Clone or download this repository.
 2. Install the [uv](https://docs.astral.sh/uv/guides/install-python/) package manager. The MCP server requires python (3.11 or later)
 3. Do a sanity check by doing
 
@@ -83,9 +83,9 @@ There are two configurations necessary before the MCP server can be invoked.
 
 ## Quick start
 
-The quickest way to do this setup is - 
+The quickest way to do this setup is -
 
-1. Create the dremio config file using 
+1. Create the dremio config file using
 
 ```shell
 $ uv run dremio-mcp-server config create dremio \
@@ -95,13 +95,13 @@ $ uv run dremio-mcp-server config create dremio \
     # --project-id <dremio project id>
 ```
 
-2. Download and install Claude desktop. And then create the Claude config file using 
+2. Download and install Claude desktop. And then create the Claude config file using
 
 ```shell
 $ uv run dremio-mcp-server config create claude
 ```
 
-3. Validate the config files using 
+3. Validate the config files using
 
 ```shell
 $ uv run dremio-mcp-server config list --type claude`
@@ -229,6 +229,18 @@ This will pickup the default location of MCP server config file. It can also be
 
 This repository is intended to be open source software that encourages contributions of any kind, like adding features, reporting issues and contributing fixes. This is not a part of Dremio product support.
 
+## Testing
+
+The project uses pytest for testing. To run the tests:
+
+```shell
+# Run all tests
+$ uv run pytest tests
+
+```
+
+GitHub Actions automatically runs tests on pull requests and pushes to the main branch.
+
 ## Contributing
 
 Please see our [Contributing Guide](CONTRIBUTING.md) for details on:
@@ -237,3 +249,4 @@ Please see our [Contributing Guide](CONTRIBUTING.md) for details on:
 -   Making contributions
 -   Code style guidelines
 -   Documentation requirements
+-   Running tests

pyproject.toml

@@ -24,6 +24,8 @@ dependencies = [
     "prompt-toolkit>=3.0.50",
     "pydantic>=2.10.6",
     "pydantic-settings>=2.8.1",
+    "pytest>=8.3.5",
+    "pytest-asyncio>=0.26.0",
     "pyyaml>=6.0.2",
     "requests>=2.32.3",
     "rich>=13.9.4",

pytest.ini

@@ -0,0 +1,13 @@
+[pytest]
+testpaths = tests
+python_files = test_*.py
+python_classes = Test*
+python_functions = test_*
+asyncio_mode = strict
+
+# Display summary info for skipped, xfailed, xpassed tests
+# along with the percentage of passing tests at the end
+addopts = -v --showlocals
+
+# Set the default timeout for async tests (in seconds)
+asyncio_timeout = 30

src/dremioai/config/settings.py

@@ -20,7 +20,7 @@
 from dremioai.config.tools import ToolType
 from enum import auto, StrEnum
 from pathlib import Path
-from yaml import safe_load
+from yaml import safe_load, add_representer, dump
 from functools import reduce
 from operator import ior
 from shutil import which
@@ -150,8 +150,8 @@ class Anthropic(BaseModel):
 
 
 class BeeAI(BaseModel):
-    mcp_server: Optional[MCPServer] = Field(default=None, alias="mcpServer")
-    sliding_memory_size: Optional[int] = Field(default=10, alias="slidingMemorySize")
+    mcp_server: Optional[MCPServer] = Field(default=None)
+    sliding_memory_size: Optional[int] = Field(default=10)
     anthropic: Optional[Anthropic] = Field(default=None)
     openai: Optional[OpenAi] = Field(default=None)
     ollama: Optional[Ollama] = Field(default=None)
@@ -223,9 +223,14 @@ def configure(cfg: Union[str, Path] = None, force=False) -> ContextVar[Settings]
     if cfg is None:
         cfg = default_config()
 
+    if not cfg.exists():
+        print(f"Creating default config file: {cfg!s}")
+        cfg.parent.mkdir(parents=True, exist_ok=True)
+        cfg.touch()
+
     with cfg.open() as f:
         s = safe_load(f)
-        _settings.set(Settings.model_validate(s))
+        _settings.set(Settings.model_validate(s if s else {}))
 
     return _settings
 
@@ -261,3 +266,31 @@ async def _call():
 
     ctx = copy_context()
     return await _call()
+
+
+def write_settings(
+    cfg: Path = None, inst: Settings = None, dry_run: bool = False
+) -> str | None:
+    if cfg is None:
+        cfg = default_config()
+
+    if not isinstance(inst, Settings):
+        inst = instance()
+
+    d = inst.model_dump(
+        exclude_none=True, mode="json", exclude_unset=True, by_alias=True
+    )
+    add_representer(
+        str,
+        lambda dumper, data: dumper.represent_scalar(
+            "tag:yaml.org,2002:str", data, style=('"' if "@" in data else None)
+        ),
+    )
+    if dry_run:
+        return dump(d)
+
+    if not cfg.exists() or not cfg.parent.exists():
+        cfg.parent.mkdir(parents=True, exist_ok=True)
+
+    with cfg.open("w") as f:
+        dump(d, f)

src/dremioai/servers/mcp.py

@@ -17,10 +17,11 @@
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.prompts import Prompt
 from mcp.server.fastmcp.resources import FunctionResource
+from mcp.cli.claude import get_claude_config_path
 from pydantic.networks import AnyUrl
 from dremioai.tools import tools
 import os
-from typing import List, Union, Annotated, Optional, Tuple
+from typing import List, Union, Annotated, Optional, Tuple, Dict, Any
 from functools import reduce
 from operator import ior
 from pathlib import Path
@@ -34,6 +35,7 @@
 from shutil import which
 import asyncio
 from yaml import dump, add_representer
+import sys
 
 
 def init(
@@ -71,10 +73,10 @@ def init(
 
 
 app = None
-if __name__ != "__main__":
-    if mode := os.environ.get("MODE"):
-        mode = [tools.ToolType[m.upper()] for m in ",".split(mode)]
-    app = init(mode=mode)
+# if __name__ != "__main__":
+# if mode := os.environ.get("MODE"):
+# mode = [tools.ToolType[m.upper()] for m in ",".split(mode)]
+# app = init(mode=mode)
 
 
 def _mode() -> List[str]:
@@ -154,6 +156,18 @@ class ConfigTypes(StrEnum):
     claude = auto()
 
 
+def get_claude_config_path() -> Path:
+    # copy of the function from mcp sdk, but returns the path whether or not
+    # it exists
+    dir = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"), "Claude")
+    match sys.platform:
+        case "win32":
+            dir = Path(Path.home(), "AppData", "Roaming", "Claude")
+        case "darwin":
+            dir = Path(Path.home(), "Library", "Application Support", "Claude")
+    return dir / "claude_desktop_config.json"
+
+
 @tc.command("list", help="Show default configuration, if it exists")
 def show_default_config(
     show_filename: Annotated[
@@ -179,13 +193,7 @@ def show_default_config(
                     )
                 )
         case ConfigTypes.claude:
-            cc = (
-                Path.home()
-                / "Library"
-                / "Application Support"
-                / "Claude"
-                / "claude_desktop_config.json"
-            )
+            cc = get_claude_config_path()
             pp(f"Default config file: '{cc!s}' (exists = {cc.exists()!s})")
             if not show_filename:
                 pp(load(cc.open()))
@@ -199,42 +207,40 @@ def show_default_config(
 tc.add_typer(cc)
 
 
-@cc.command("claude", help="Create a default configuration file for Claude")
-def create_default_config(
-    dry_run: Annotated[
-        bool, Option(help="Dry run, do not overwrite the config file. Just print it")
-    ] = False,
-):
+def create_default_mcpserver_config() -> Dict[str, Any]:
     if (uv := which("uv")) is not None:
         uv = Path(uv).resolve()
         dir = str(Path(os.getcwd()).resolve())
-        dmcp = {
-            "Dremio": {
-                "command": str(uv),
-                "args": ["run", "--directory", dir, "dremio-mcp-server", "run"],
-            }
+        return {
+            "command": str(uv),
+            "args": ["run", "--directory", dir, "dremio-mcp-server", "run"],
         }
-        cc = (
-            Path.home()
-            / "Library"
-            / "Application Support"
-            / "Claude"
-            / "claude_desktop_config.json"
-        )
-        c = load(cc.open()) if cc.exists() else {"mcpServers": {}}
-        c["mcpServers"].update(dmcp)
-        if dry_run:
-            pp(c)
-        else:
-            if not cc.exists():
-                cc.parent.mkdir(parents=True, exist_ok=True)
-            with cc.open("w") as f:
-                jdump(c, f)
-                pp(f"Created default config file: {cc!s}")
     else:
         raise FileNotFoundError("uv command not found. Please install uv")
 
 
+@cc.command("claude", help="Create a default configuration file for Claude")
+def create_default_config(
+    dry_run: Annotated[
+        bool, Option(help="Dry run, do not overwrite the config file. Just print it")
+    ] = False,
+):
+    cc = get_claude_config_path()
+    dcmp = {"Dremio": cc}
+    c = load(cc.open()) if cc.exists() else {"mcpServers": {}}
+    c["mcpServers"].update(dcmp)
+    if dry_run:
+        pp(c)
+        return
+
+    if not cc.exists():
+        cc.parent.mkdir(parents=True, exist_ok=True)
+
+    with cc.open("w") as f:
+        jdump(c, f)
+        pp(f"Created default config file: {cc!s}")
+
+
 @cc.command("dremioai", help="Create a default configuration file")
 def create_default_config(
     uri: Annotated[
@@ -273,7 +279,7 @@ def create_default_config(
             "enable_experimental": enable_experimental,
         }
     )
-    ts = settings.Tools.model_validate({"server_mode": tools.ToolType.FOR_SELF})
+    ts = settings.Tools.model_validate({"server_mode": mode})
     settings.configure(settings.default_config(), force=True)
     settings.instance().dremio = dremio
     settings.instance().tools = ts