test_ai_command_utils.py
1 import platform 2 from unittest import mock 3 4 import pytest 5 6 from mlflow.ai_commands import get_command, get_command_body, list_commands, parse_frontmatter 7 8 9 def test_parse_frontmatter_with_metadata(): 10 content = """--- 11 namespace: genai 12 description: Test command 13 --- 14 15 # Command content 16 This is the body.""" 17 18 metadata, body = parse_frontmatter(content) 19 20 assert metadata["namespace"] == "genai" 21 assert metadata["description"] == "Test command" 22 assert "# Command content" in body 23 assert "This is the body." in body 24 25 26 def test_parse_frontmatter_without_metadata(): 27 content = "# Just a regular markdown file\nNo frontmatter here." 28 29 metadata, body = parse_frontmatter(content) 30 31 assert metadata == {} 32 assert body == content 33 34 35 def test_parse_frontmatter_malformed(): 36 content = """--- 37 invalid: yaml: [ 38 --- 39 Body content""" 40 41 # Should not raise, but return empty metadata 42 metadata, body = parse_frontmatter(content) 43 assert metadata == {} 44 assert body == content 45 46 47 def test_parse_frontmatter_empty_metadata(): 48 content = """--- 49 --- 50 Body content""" 51 52 metadata, body = parse_frontmatter(content) 53 # Empty YAML returns None, which becomes {} 54 assert metadata == {} or metadata is None 55 assert "Body content" in body 56 57 58 def test_list_commands_all(tmp_path): 59 # Create test command structure 60 genai_dir = tmp_path / "commands" / "genai" 61 genai_dir.mkdir(parents=True) 62 63 test_cmd = genai_dir / "test.md" 64 test_cmd.write_text("""--- 65 namespace: genai 66 description: Test command 67 --- 68 Content""") 69 70 another_dir = tmp_path / "commands" / "ml" 71 another_dir.mkdir(parents=True) 72 73 another_cmd = another_dir / "train.md" 74 another_cmd.write_text("""--- 75 namespace: ml 76 description: Training command 77 --- 78 Content""") 79 80 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 81 # Mock Path(__file__).parent to return tmp_path/commands 82 mock_path.return_value.parent = tmp_path / "commands" 83 84 commands = list_commands() 85 86 assert len(commands) == 2 87 # Use forward slashes consistently in assertions 88 assert any(cmd["key"] == "genai/test" for cmd in commands) 89 assert any(cmd["key"] == "ml/train" for cmd in commands) 90 91 92 def test_list_commands_with_namespace_filter(tmp_path): 93 # Setup test commands 94 genai_dir = tmp_path / "commands" / "genai" 95 genai_dir.mkdir(parents=True) 96 97 cmd1 = genai_dir / "analyze.md" 98 cmd1.write_text("""--- 99 namespace: genai 100 description: Analyze command 101 --- 102 Content""") 103 104 cmd2 = genai_dir / "evaluate.md" 105 cmd2.write_text("""--- 106 namespace: genai 107 description: Evaluate command 108 --- 109 Content""") 110 111 ml_dir = tmp_path / "commands" / "ml" 112 ml_dir.mkdir(parents=True) 113 114 cmd3 = ml_dir / "train.md" 115 cmd3.write_text("""--- 116 namespace: ml 117 description: Training command 118 --- 119 Content""") 120 121 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 122 mock_path.return_value.parent = tmp_path / "commands" 123 124 # Filter by genai namespace 125 genai_commands = list_commands(namespace="genai") 126 127 assert len(genai_commands) == 2 128 assert all(cmd["key"].startswith("genai/") for cmd in genai_commands) 129 130 131 def test_get_command_success(tmp_path): 132 genai_dir = tmp_path / "commands" / "genai" 133 genai_dir.mkdir(parents=True) 134 135 test_content = """--- 136 namespace: genai 137 description: Test command 138 --- 139 140 # Test Command 141 This is the full content.""" 142 143 test_cmd = genai_dir / "analyze.md" 144 test_cmd.write_text(test_content) 145 146 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 147 mock_path.return_value.parent = tmp_path / "commands" 148 149 content = get_command("genai/analyze") 150 151 assert content == test_content 152 153 154 def test_get_command_not_found(tmp_path): 155 commands_dir = tmp_path / "commands" 156 commands_dir.mkdir() 157 158 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 159 mock_path.return_value.parent = commands_dir 160 161 with pytest.raises(FileNotFoundError, match="Command 'nonexistent/command' not found"): 162 get_command("nonexistent/command") 163 164 165 def test_list_commands_empty_directory(tmp_path): 166 # Create empty commands directory 167 commands_dir = tmp_path / "commands" 168 commands_dir.mkdir() 169 170 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 171 mock_path.return_value.parent = tmp_path 172 173 commands = list_commands() 174 175 assert commands == [] 176 177 178 def test_list_commands_nonexistent_directory(tmp_path): 179 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 180 mock_path.return_value.parent = tmp_path 181 182 commands = list_commands() 183 184 assert commands == [] 185 186 187 def test_list_commands_with_invalid_files(tmp_path): 188 genai_dir = tmp_path / "commands" / "genai" 189 genai_dir.mkdir(parents=True) 190 191 # Valid command 192 valid_cmd = genai_dir / "valid.md" 193 valid_cmd.write_text("""--- 194 namespace: genai 195 description: Valid command 196 --- 197 Content""") 198 199 # Create a file with invalid YAML to trigger parsing error 200 invalid_cmd = genai_dir / "invalid.md" 201 invalid_cmd.write_text("Invalid content that will cause parsing error") 202 203 # On Unix-like systems, remove read permissions 204 if platform.system() != "Windows": 205 invalid_cmd.chmod(0o000) 206 207 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 208 mock_path.return_value.parent = tmp_path / "commands" 209 210 commands = list_commands() 211 212 # Restore permissions for cleanup 213 if platform.system() != "Windows": 214 invalid_cmd.chmod(0o644) 215 216 # Should include both commands (invalid one gets parsed but with empty metadata) 217 assert len(commands) >= 1 218 # Ensure we have at least the valid command 219 valid_commands = [cmd for cmd in commands if cmd["key"] == "genai/valid"] 220 assert len(valid_commands) == 1 221 assert valid_commands[0]["description"] == "Valid command" 222 223 224 def test_list_commands_sorted(): 225 # Use the real implementation with actual files 226 commands = list_commands() 227 228 # If there are any commands, verify they're sorted 229 if len(commands) > 1: 230 keys = [cmd["key"] for cmd in commands] 231 assert keys == sorted(keys) 232 233 234 def test_get_command_body(tmp_path): 235 genai_dir = tmp_path / "commands" / "genai" 236 genai_dir.mkdir(parents=True) 237 238 # Test with frontmatter 239 content_with_frontmatter = """--- 240 namespace: genai 241 description: Test command 242 --- 243 244 # Test Command 245 This is the body content.""" 246 247 test_cmd = genai_dir / "analyze.md" 248 test_cmd.write_text(content_with_frontmatter) 249 250 # Test without frontmatter - should return entire content 251 content_no_frontmatter = """# Simple Command 252 This is just markdown content.""" 253 254 simple_cmd = genai_dir / "simple.md" 255 simple_cmd.write_text(content_no_frontmatter) 256 257 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 258 mock_path.return_value.parent = tmp_path / "commands" 259 260 # Test with frontmatter 261 body = get_command_body("genai/analyze") 262 263 # Should strip frontmatter and return only body 264 assert "namespace: genai" not in body 265 assert "description: Test command" not in body 266 assert "# Test Command" in body 267 assert "This is the body content." in body 268 269 # Test without frontmatter 270 body_no_frontmatter = get_command_body("genai/simple") 271 assert body_no_frontmatter == content_no_frontmatter 272 273 274 def test_get_command_body_not_found(tmp_path): 275 commands_dir = tmp_path / "commands" 276 commands_dir.mkdir() 277 278 with mock.patch("mlflow.ai_commands.ai_command_utils.Path") as mock_path: 279 mock_path.return_value.parent = commands_dir 280 281 with pytest.raises(FileNotFoundError, match="Command 'nonexistent/command' not found"): 282 get_command_body("nonexistent/command")