import unittest
from unittest.mock import patch, MagicMock
from armorer.agent import shell

class TestShellTool(unittest.TestCase):
    
    @patch('subprocess.run')
    def test_shell_truncates_long_output(self, mock_run):
        # Create a mock result with very long stdout
        long_output = "a" * 25000
        mock_result = MagicMock()
        mock_result.stdout = long_output
        mock_result.stderr = ""
        mock_result.returncode = 0
        mock_run.return_value = mock_result
        
        # Execute shell tool
        result = shell("echo lengthy", timeout=1)
        
        # Verify length
        # Should be 20000 + warning message length
        self.assertLess(len(result), 25000)
        self.assertTrue(result.endswith("]"))
        self.assertIn("... [Output truncated.", result)
        
        # Verify content
        self.assertTrue(result.startswith("aaaaa"))
        
    @patch('subprocess.run')
    def test_shell_keeps_short_output(self, mock_run):
        # Create a mock result with short stdout
        short_output = "short output"
        mock_result = MagicMock()
        mock_result.stdout = short_output
        mock_result.stderr = ""
        mock_result.returncode = 0
        mock_run.return_value = mock_result
        
        # Execute shell tool
        result = shell("echo short", timeout=1)
        
        # Verify content
        self.assertEqual(result, short_output)
        
if __name__ == '__main__':
    unittest.main()
