1 line
9.7 KiB
JSON
1 line
9.7 KiB
JSON
{"repo": "calinfaja/k-lean", "files": {"tests/unit/test_model_management.py": {"c8a06e9c161132278231bb2fca91d86773f659e4": {"t": 1785596049.8191385, "n": 9151, "h": "1bda2a4c2f156e26", "c": "\"\"\"Tests for model management (model add, model remove commands).\"\"\"\n\nimport tempfile\nfrom pathlib import Path\nfrom unittest.mock import patch\n\nfrom click.testing import CliRunner\n\nfrom klean.cli import main\nfrom klean.model_utils import extract_model_name, is_thinking_model, parse_model_id\n\n\nclass TestModelUtils:\n \"\"\"Tests for model utility functions.\"\"\"\n\n def test_extract_model_name_openrouter(self):\n \"\"\"Test extracting model name from OpenRouter ID.\"\"\"\n result = extract_model_name(\"openrouter/kwaipilot/kat-coder-pro:free\")\n assert result == \"kat-coder-pro\"\n\n def test_extract_model_name_nanogpt(self):\n \"\"\"Test extracting model name from NanoGPT ID.\"\"\"\n result = extract_model_name(\"openai/Qwen/Qwen2.5-Coder-32B-Instruct\")\n # Should extract meaningful part (converts to lowercase, removes size info)\n assert \"qwen\" in result.lower()\n assert len(result) > 3, \"Model name should not be too abbreviated\"\n\n def test_extract_model_name_simple(self):\n \"\"\"Test extracting model name from simple ID.\"\"\"\n result = extract_model_name(\"anthropic/claude-3.5-sonnet\")\n assert \"claude\" in result\n assert \"sonnet\" in result\n\n def test_extract_model_name_with_version(self):\n \"\"\"Test extracting model name with version numbers.\"\"\"\n result = extract_model_name(\"openai/Qwen/Qwen2.5-Coder\")\n # Should replace dots with dashes\n assert \".\" not in result\n assert \"-\" in result\n\n def test_parse_model_id_openrouter(self):\n \"\"\"Test parsing OpenRouter model ID.\"\"\"\n provider, full_id = parse_model_id(\"openrouter/anthropic/claude-3.5-sonnet\")\n assert provider == \"openrouter\"\n assert full_id == \"openrouter/anthropic/claude-3.5-sonnet\"\n\n def test_parse_model_id_nanogpt(self):\n \"\"\"Test parsing NanoGPT (OpenAI) model ID.\"\"\"\n provider, full_id = parse_model_id(\"openai/Qwen/Qwen2.5-Coder-32B-Instruct\")\n assert provider == \"nanogpt\"\n assert full_id == \"openai/Qwen/Qwen2.5-Coder-32B-Instruct\"\n\n def test_parse_model_id_unknown(self):\n \"\"\"Test parsing unknown provider ID.\"\"\"\n provider, full_id = parse_model_id(\"unknown/model\")\n assert provider == \"unknown\"\n\n def test_is_thinking_model_true(self):\n \"\"\"Test thinking model detection.\"\"\"\n assert is_thinking_model(\"deepseek-r1-thinking\")\n assert is_thinking_model(\"glm-4.6-thinking\")\n assert is_thinking_model(\"kimi-k2-thinking\")\n\n def test_is_thinking_model_false(self):\n \"\"\"Test non-thinking model detection.\"\"\"\n assert not is_thinking_model(\"claude-3.5-sonnet\")\n assert not is_thinking_model(\"qwen-coder\")\n assert not is_thinking_model(\"deepseek-v3\")\n\n\nclass TestAddModelCommand:\n \"\"\"Tests for kln model add command.\"\"\"\n\n def test_add_model_command_exists(self):\n \"\"\"Test that kln model add command exists.\"\"\"\n runner = CliRunner()\n result = runner.invoke(main, [\"model\", \"add\", \"--help\"])\n assert result.exit_code == 0\n assert \"Add individual model\" in result.output\n\n def test_add_model_no_config(self):\n \"\"\"Test add-model fails when config doesn't exist.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n result = runner.invoke(\n main,\n [\"model\", \"add\", \"--provider\", \"openrouter\", \"anthropic/claude-3.5-sonnet\"],\n )\n\n assert result.exit_code != 0\n assert \"config.yaml not found\" in result.output\n\n def test_add_model_openrouter(self):\n \"\"\"Test adding OpenRouter model.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n config_dir.mkdir(parents=True, exist_ok=True)\n\n # Create initial config\n config_yaml = \"\"\"litellm_settings:\n drop_params: true\n\nmodel_list:\n - model_name: test-model\n litellm_params:\n model: openai/test\n api_key: os.environ/NANOGPT_API_KEY\n\"\"\"\n (config_dir / \"config.yaml\").write_text(config_yaml)\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n result = runner.invoke(\n main,\n [\"model\", \"add\", \"--provider\", \"openrouter\", \"anthropic/claude-3.5-sonnet\"],\n )\n\n assert result.exit_code == 0\n assert \"Added model\" in result.output\n\n # Verify model was added\n config_content = (config_dir / \"config.yaml\").read_text()\n assert \"claude-3.5-sonnet\" in config_content\n assert \"openrouter/anthropic\" in config_content\n\n\nclass TestRemoveModelCommand:\n \"\"\"Tests for kln model remove command.\"\"\"\n\n def test_remove_model_command_exists(self):\n \"\"\"Test that kln model remove command exists.\"\"\"\n runner = CliRunner()\n result = runner.invoke(main, [\"model\", \"remove\", \"--help\"])\n assert result.exit_code == 0\n assert \"Remove model\" in result.output\n\n def test_remove_model_no_config(self):\n \"\"\"Test remove-model fails when config doesn't exist.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n result = runner.invoke(main, [\"model\", \"remove\", \"nonexistent-model\"])\n\n assert result.exit_code != 0\n assert \"config.yaml not found\" in result.output\n\n def test_remove_model_not_found(self):\n \"\"\"Test remove-model when model doesn't exist.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n config_dir.mkdir(parents=True, exist_ok=True)\n\n # Create config with one model\n config_yaml = \"\"\"litellm_settings:\n drop_params: true\n\nmodel_list:\n - model_name: test-model\n litellm_params:\n model: openai/test\n api_key: os.environ/NANOGPT_API_KEY\n\"\"\"\n (config_dir / \"config.yaml\").write_text(config_yaml)\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n result = runner.invoke(main, [\"model\", \"remove\", \"nonexistent-model\"])\n\n assert result.exit_code == 0\n assert \"not found\" in result.output\n assert \"Available models:\" in result.output\n\n def test_remove_model_success(self):\n \"\"\"Test successfully removing a model.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n config_dir.mkdir(parents=True, exist_ok=True)\n\n # Create config with two models\n config_yaml = \"\"\"litellm_settings:\n drop_params: true\n\nmodel_list:\n - model_name: test-model-1\n litellm_params:\n model: openai/test1\n api_key: os.environ/NANOGPT_API_KEY\n - model_name: test-model-2\n litellm_params:\n model: openai/test2\n api_key: os.environ/NANOGPT_API_KEY\n\"\"\"\n (config_dir / \"config.yaml\").write_text(config_yaml)\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n result = runner.invoke(main, [\"model\", \"remove\", \"test-model-1\"])\n\n assert result.exit_code == 0\n assert \"Removed model\" in result.output\n\n # Verify model was removed\n config_content = (config_dir / \"config.yaml\").read_text()\n assert \"test-model-1\" not in config_content\n assert \"test-model-2\" in config_content\n\n\nclass TestAddAndRemoveIntegration:\n \"\"\"Integration tests for add and remove operations.\"\"\"\n\n def test_add_then_remove_model(self):\n \"\"\"Test adding then removing a model.\"\"\"\n runner = CliRunner()\n\n with tempfile.TemporaryDirectory() as tmpdir:\n config_dir = Path(tmpdir) / \".config\" / \"litellm\"\n config_dir.mkdir(parents=True, exist_ok=True)\n\n # Create initial config\n config_yaml = \"\"\"litellm_settings:\n drop_params: true\n\nmodel_list: []\n\"\"\"\n (config_dir / \"config.yaml\").write_text(config_yaml)\n\n with patch(\"klean.cli.CONFIG_DIR\", config_dir):\n # Add model\n result1 = runner.invoke(\n main,\n [\"model\", \"add\", \"--provider\", \"openrouter\", \"anthropic/claude-3.5-sonnet\"],\n )\n assert result1.exit_code == 0\n\n # Remove model\n result2 = runner.invoke(main, [\"model\", \"remove\", \"claude-3.5-sonnet\"])\n assert result2.exit_code == 0\n\n # Verify model list is empty\n config_content = (config_dir / \"config.yaml\").read_text()\n assert \"claude-3.5-sonnet\" not in config_content\n"}}}} |