|
4 | 4 | """
|
5 | 5 | import csv
|
6 | 6 | import os
|
| 7 | +from unittest.mock import AsyncMock, patch |
7 | 8 |
|
8 | 9 | import pytest
|
9 | 10 |
|
10 | 11 | from codegate.pipeline.suspicious_commands.suspicious_commands import (
|
11 | 12 | SuspiciousCommands,
|
| 13 | + check_suspicious_code, |
12 | 14 | )
|
13 | 15 |
|
14 | 16 | try:
|
@@ -189,3 +191,68 @@ async def test_classify_phrase_confident(sc):
|
189 | 191 | else:
|
190 | 192 | print(f"{command['cmd']} {prob} {prediction} 1")
|
191 | 193 | check_results(tp, tn, fp, fn)
|
| 194 | + |
| 195 | + |
| 196 | +@pytest.mark.asyncio |
| 197 | +@patch("codegate.pipeline.suspicious_commands.suspicious_commands.SuspiciousCommands.get_instance") |
| 198 | +async def test_check_suspicious_code_safe(mock_get_instance): |
| 199 | + """ |
| 200 | + Test check_suspicious_code with safe code. |
| 201 | + """ |
| 202 | + mock_instance = mock_get_instance.return_value |
| 203 | + mock_instance.classify_phrase = AsyncMock(return_value=(0, 0.5)) |
| 204 | + |
| 205 | + code = "print('Hello, world!')" |
| 206 | + comment, is_suspicious = await check_suspicious_code(code, "python") |
| 207 | + |
| 208 | + assert comment == "" |
| 209 | + assert is_suspicious is False |
| 210 | + |
| 211 | + |
| 212 | +@pytest.mark.asyncio |
| 213 | +@patch("codegate.pipeline.suspicious_commands.suspicious_commands.SuspiciousCommands.get_instance") |
| 214 | +async def test_check_suspicious_code_suspicious(mock_get_instance): |
| 215 | + """ |
| 216 | + Test check_suspicious_code with suspicious code. |
| 217 | + """ |
| 218 | + mock_instance = mock_get_instance.return_value |
| 219 | + mock_instance.classify_phrase = AsyncMock(return_value=(1, 0.95)) |
| 220 | + |
| 221 | + code = "rm -rf /" |
| 222 | + comment, is_suspicious = await check_suspicious_code(code, "bash") |
| 223 | + |
| 224 | + assert "🛡️ CodeGate: The bash supplied is likely unsafe." in comment |
| 225 | + assert is_suspicious is True |
| 226 | + |
| 227 | + |
| 228 | +@pytest.mark.asyncio |
| 229 | +@patch("codegate.pipeline.suspicious_commands.suspicious_commands.SuspiciousCommands.get_instance") |
| 230 | +async def test_check_suspicious_code_skipped_language(mock_get_instance): |
| 231 | + """ |
| 232 | + Test check_suspicious_code with a language that should be skipped. |
| 233 | + """ |
| 234 | + mock_instance = mock_get_instance.return_value |
| 235 | + mock_instance.classify_phrase = AsyncMock() |
| 236 | + |
| 237 | + code = "print('Hello, world!')" |
| 238 | + comment, is_suspicious = await check_suspicious_code(code, "python") |
| 239 | + |
| 240 | + assert comment == "" |
| 241 | + assert is_suspicious is False |
| 242 | + mock_instance.classify_phrase.assert_not_called() |
| 243 | + |
| 244 | + |
| 245 | +@pytest.mark.asyncio |
| 246 | +@patch("codegate.pipeline.suspicious_commands.suspicious_commands.SuspiciousCommands.get_instance") |
| 247 | +async def test_check_suspicious_code_no_language(mock_get_instance): |
| 248 | + """ |
| 249 | + Test check_suspicious_code with no language specified. |
| 250 | + """ |
| 251 | + mock_instance = mock_get_instance.return_value |
| 252 | + mock_instance.classify_phrase = AsyncMock(return_value=(1, 0.85)) |
| 253 | + |
| 254 | + code = "rm -rf /" |
| 255 | + comment, is_suspicious = await check_suspicious_code(code) |
| 256 | + |
| 257 | + assert "🛡️ CodeGate: The code supplied is possibly unsafe." in comment |
| 258 | + assert is_suspicious is True |
0 commit comments