zk-data-agent/tests/test_bundled_skills.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from src.bundled_skills import (
    find_bundled_skill,
    format_skills_for_system_prompt,
    get_bundled_skills,
    load_directory_skills,
    load_project_skills,
)
from src.agent_runtime import LocalCodingAgent
from src.agent_types import AgentRuntimeConfig, ModelConfig


class BundledSkillsTests(unittest.TestCase):
    def test_directory_skills_are_loaded_with_metadata(self) -> None:
        skills = {skill.name: skill for skill in get_bundled_skills()}
        self.assertIn('verify', skills)
        self.assertEqual(skills['verify'].source, 'directory')
        self.assertIn('bash', skills['verify'].allowed_tools)

    def test_directory_skill_aliases_are_resolved(self) -> None:
        skill = find_bundled_skill('config-help')
        self.assertIsNotNone(skill)
        assert skill is not None
        self.assertEqual(skill.name, 'update-config')
        self.assertEqual(skill.source, 'directory')

    def test_directory_skill_prompt_appends_invocation_args(self) -> None:
        skill = find_bundled_skill('verify')
        self.assertIsNotNone(skill)
        assert skill is not None
        prompt = skill.get_prompt(None, 'focus: tests')  # type: ignore[arg-type]
        self.assertIn('Verify that the recent code changes work correctly.', prompt)
        self.assertIn('## Invocation Arguments', prompt)
        self.assertIn('focus: tests', prompt)

    def test_custom_directory_skill_loader(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            skill_dir = root / 'sample'
            skill_dir.mkdir()
            (skill_dir / 'SKILL.md').write_text(
                (
                    '---\n'
                    'name: sample\n'
                    'description: Sample directory skill.\n'
                    'aliases: s1, s2\n'
                    'allowed_tools: read_file, write_file\n'
                    '---\n'
                    'Use sample instructions.\n'
                ),
                encoding='utf-8',
            )
            skills = load_directory_skills(root)
        self.assertEqual(len(skills), 1)
        self.assertEqual(skills[0].name, 'sample')
        self.assertEqual(skills[0].aliases, ('s1', 's2'))
        self.assertEqual(skills[0].allowed_tools, ('read_file', 'write_file'))

    def test_skills_prompt_includes_directory_skill(self) -> None:
        rendered = format_skills_for_system_prompt()
        self.assertIn('verify', rendered)
        self.assertIn('update-config', rendered)

    def test_skills_prompt_respects_enabled_skill_filter(self) -> None:
        rendered = format_skills_for_system_prompt(
            enabled_skill_names=('verify', 'product-data')
        )
        self.assertIn('verify', rendered)
        self.assertIn('product-data', rendered)
        self.assertNotIn('online-mining', rendered)

    def test_project_skills_are_loaded_from_workspace_skills_directory(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            skill_dir = workspace / 'skills' / 'project-skill'
            skill_dir.mkdir(parents=True)
            (skill_dir / 'SKILL.md').write_text(
                (
                    '---\n'
                    'name: project-skill\n'
                    'description: Project-maintained skill.\n'
                    'aliases: local-skill\n'
                    'allowed_tools: read_file\n'
                    '---\n'
                    'Project skill body.\n'
                ),
                encoding='utf-8',
            )
            skills = get_bundled_skills(workspace)
            resolved = find_bundled_skill('local-skill', cwd=workspace)
        by_name = {skill.name: skill for skill in skills}
        self.assertIn('project-skill', by_name)
        self.assertEqual(by_name['project-skill'].source, 'project')
        self.assertIsNotNone(resolved)
        assert resolved is not None
        self.assertEqual(resolved.name, 'project-skill')

    def test_project_skills_override_bundled_directory_skills(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            skill_dir = workspace / 'skills' / 'verify'
            skill_dir.mkdir(parents=True)
            (skill_dir / 'SKILL.md').write_text(
                (
                    '---\n'
                    'name: verify\n'
                    'description: Project override for verify.\n'
                    '---\n'
                    'Project verify body.\n'
                ),
                encoding='utf-8',
            )
            skill = find_bundled_skill('verify', cwd=workspace)
        self.assertIsNotNone(skill)
        assert skill is not None
        self.assertEqual(skill.source, 'project')
        self.assertIn('Project verify body.', skill.get_prompt(None, ''))  # type: ignore[arg-type]

    def test_load_project_skills_returns_empty_without_skills_directory(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            self.assertEqual(load_project_skills(Path(tmp_dir)), ())

    def test_agent_skill_tool_executes_project_skill_from_cwd(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            workspace = Path(tmp_dir)
            skill_dir = workspace / 'skills' / 'project-run'
            skill_dir.mkdir(parents=True)
            (skill_dir / 'SKILL.md').write_text(
                (
                    '---\n'
                    'name: project-run\n'
                    'description: Project runnable skill.\n'
                    '---\n'
                    'Run project skill body.\n'
                ),
                encoding='utf-8',
            )
            agent = LocalCodingAgent(
                model_config=ModelConfig(model='test'),
                runtime_config=AgentRuntimeConfig(cwd=workspace),
            )
            result = agent._execute_skill({'skill': 'project-run', 'args': 'demo'})
        self.assertTrue(result.ok)
        self.assertIn('Run project skill body.', result.content)
        self.assertEqual(result.metadata.get('source'), 'project')
        self.assertIn('SKILL.md', str(result.metadata.get('skill_path')))


if __name__ == '__main__':
    unittest.main()