""" Unit Tests for inference.py ============================= Tests LLM interaction, parsing, retry logic, context window, and output formatting. Run: python test_inference.py [--verbose] """ import json import unittest from io import StringIO import sys from unittest.mock import Mock, patch, MagicMock # Import functions to test from inference import ( parse_llm_response, create_fallback_action, ContextWindow, emit_start, emit_step, emit_end, call_llm, run_episode, MAX_TOKENS, MAX_CONTEXT_TURNS, ) class TestLLMParsing(unittest.TestCase): """Test LLM response parsing with various formats""" def test_parse_clean_json(self): """Test parsing clean JSON response""" response = json.dumps({ "thought": "Check alerts", "tool": "aws.soc.get_alerts", "args": {} }) parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIsNone(error) self.assertEqual(parsed["tool"], "aws.soc.get_alerts") self.assertEqual(parsed["thought"], "Check alerts") def test_parse_json_with_markdown_block(self): """Test parsing JSON inside markdown code block""" response = '''Here's my action: ```json { "thought": "Query logs", "tool": "aws.cloudwatch.query_deep", "args": {"log_group": "/aws/ec2"} } ``` Now I'll wait for results.''' parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIsNone(error) self.assertEqual(parsed["tool"], "aws.cloudwatch.query_deep") def test_parse_json_embedded_in_text(self): """Test parsing JSON embedded in explanation text""" response = '''I need to check the alerts. My action is: { "thought": "Get current alerts", "tool": "aws.soc.get_alerts", "args": {} } This will show me what's happening.''' parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIsNone(error) self.assertEqual(parsed["tool"], "aws.soc.get_alerts") def test_parse_json_with_single_quotes(self): """Test parsing JSON with single quotes (malformed but recoverable)""" response = "{'thought': 'Check alerts', 'tool': 'aws.soc.get_alerts', 'args': {}}" parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIsNone(error) self.assertEqual(parsed["tool"], "aws.soc.get_alerts") def test_parse_json_with_trailing_comma(self): """Test parsing JSON with trailing commas (malformed but recoverable)""" response = '''{ "thought": "Check alerts", "tool": "aws.soc.get_alerts", "args": {}, }''' parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIsNone(error) self.assertEqual(parsed["tool"], "aws.soc.get_alerts") def test_parse_missing_tool_field(self): """Test handling of JSON missing tool field""" response = json.dumps({ "thought": "Something", "args": {} }) parsed, error = parse_llm_response(response) self.assertIsNone(parsed) self.assertIsNotNone(error) self.assertIn("tool", error.lower()) def test_parse_empty_response(self): """Test handling of empty response""" parsed, error = parse_llm_response("") self.assertIsNone(parsed) self.assertIsNotNone(error) def test_parse_completely_invalid_json(self): """Test handling of completely invalid JSON""" response = "this is not json at all gdfshjkl" parsed, error = parse_llm_response(response) self.assertIsNone(parsed) self.assertIsNotNone(error) def test_parse_adds_missing_args(self): """Test that missing args dict is added""" response = json.dumps({ "thought": "Act", "tool": "aws.soc.get_alerts" }) parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIn("args", parsed) self.assertEqual(parsed["args"], {}) def test_parse_adds_default_thought(self): """Test that missing thought is added""" response = json.dumps({ "tool": "aws.soc.get_alerts", "args": {} }) parsed, error = parse_llm_response(response) self.assertIsNotNone(parsed) self.assertIn("thought", parsed) class TestFallbackAction(unittest.TestCase): """Test fallback action generation""" def test_fallback_action_structure(self): """Test that fallback action has required fields""" action = create_fallback_action(step=1) self.assertIn("tool", action) self.assertIn("args", action) self.assertIn("thought", action) def test_fallback_varies_by_step(self): """Test that fallback varies by step to avoid loops""" action1 = create_fallback_action(step=1) action2 = create_fallback_action(step=2) action3 = create_fallback_action(step=3) tools = [action1["tool"], action2["tool"], action3["tool"]] # Should not all be the same self.assertGreater(len(set(tools)), 1) def test_fallback_has_valid_tool(self): """Test that fallback action uses valid tools""" valid_tools = [ "aws.soc.get_alerts", "aws.guardduty.get_findings", "aws.cloudtrail.lookup_events", "aws.ec2.describe", "aws.iam.describe_role", ] action = create_fallback_action(step=1) self.assertIn(action["tool"], valid_tools) class TestContextWindow(unittest.TestCase): """Test context window sliding mechanism""" def test_context_window_init(self): """Test context window initialization""" context = ContextWindow(system_prompt="Test prompt", max_turns=3) self.assertEqual(context.system_prompt, "Test prompt") self.assertEqual(context.max_turns, 3) self.assertEqual(len(context.turns), 0) def test_context_window_add_turn(self): """Test adding turns to context""" context = ContextWindow(max_turns=3) context.add_turn( observation="Obs 1", action="Action 1", result="Result 1" ) self.assertEqual(len(context.turns), 1) self.assertEqual(context.turns[0]["observation"], "Obs 1") def test_context_window_sliding(self): """Test that context window slides when exceeding max""" context = ContextWindow(max_turns=2) for i in range(5): context.add_turn( observation=f"Obs {i}", action=f"Action {i}", result=f"Result {i}" ) # Should only keep last 2 turns self.assertEqual(len(context.turns), 2) self.assertEqual(context.turns[0]["observation"], "Obs 3") self.assertEqual(context.turns[1]["observation"], "Obs 4") def test_context_get_messages_format(self): """Test that get_messages returns proper OpenAI format""" context = ContextWindow(system_prompt="System", max_turns=2) context.add_turn( observation="Obs", action='{"tool": "test"}', result="Result" ) messages = context.get_messages() # Should have system + user + assistant self.assertGreaterEqual(len(messages), 1) self.assertEqual(messages[0]["role"], "system") self.assertEqual(messages[0]["content"], "System") def test_context_clear(self): """Test clearing context""" context = ContextWindow() context.add_turn("O", "A", "R") context.clear() self.assertEqual(len(context.turns), 0) class TestOutputFormatting(unittest.TestCase): """Test stdout output formatting for hackathon compliance""" def test_emit_start_format(self): """Test [START] line format""" output = StringIO() with patch('sys.stdout', output): emit_start(task_name="easy", env_name="cloudsoc", model_name="TestModel") result = output.getvalue() self.assertIn("[START]", result) self.assertIn("task=easy", result) self.assertIn("env=cloudsoc", result) self.assertIn("model=TestModel", result) def test_emit_step_format(self): """Test [STEP] line format""" output = StringIO() with patch('sys.stdout', output): emit_step( step=1, action='aws.soc.get_alerts({})', reward=0.05, done=False, error=None ) result = output.getvalue() self.assertIn("[STEP]", result) self.assertIn("step=1", result) self.assertIn("reward=0.05", result) self.assertIn("done=false", result) self.assertIn("error=null", result) def test_emit_step_with_error(self): """Test [STEP] line with error""" output = StringIO() with patch('sys.stdout', output): emit_step( step=1, action='aws.invalid()', reward=-0.02, done=False, error="VALIDATION_ERROR" ) result = output.getvalue() self.assertIn("error=VALIDATION_ERROR", result) def test_emit_end_format(self): """Test [END] line format""" output = StringIO() with patch('sys.stdout', output): emit_end(success=True, steps=5, score=0.75, rewards=[0.1, 0.2, 0.15, 0.3, 0.0]) result = output.getvalue() self.assertIn("[END]", result) self.assertIn("success=true", result) self.assertIn("steps=5", result) # Score should be between 0.1 and 0.9 self.assertIn("score=0.750", result) score_val = float(result.split('score=')[1].split()[0]) self.assertGreaterEqual(score_val, 0.1) self.assertLessEqual(score_val, 0.9) self.assertIn("rewards=0.10,0.20,0.15,0.30,0.00", result) def test_emit_end_newlines(self): """Test that output lines have no embedded newlines""" output = StringIO() with patch('sys.stdout', output): emit_step( step=1, action='tool({"param": "value"})', reward=0.0, done=False, error=None ) result = output.getvalue().strip() # Should be single line self.assertEqual(result.count('\n'), 0) class TestCallLLMFallback(unittest.TestCase): """Test LLM calling with error handling""" @patch('inference.client') def test_call_llm_success(self, mock_client): """Test successful LLM call""" mock_response = MagicMock() mock_response.choices[0].message.content = '{"tool": "test", "args": {}}' mock_client.chat.completions.create.return_value = mock_response result = call_llm([{"role": "user", "content": "Test"}]) self.assertIsNotNone(result) self.assertIn("tool", result) @patch('inference.client') def test_call_llm_timeout_fallback(self, mock_client): """Test that LLM timeout returns fallback""" mock_client.chat.completions.create.side_effect = TimeoutError("Timeout") result = call_llm([{"role": "user", "content": "Test"}]) # Should return fallback action string self.assertIsNotNone(result) self.assertIsInstance(result, str) class TestMaxConstraints(unittest.TestCase): """Test that 2vCPU/8GB constraints are respected""" def test_max_tokens_set(self): """Test that MAX_TOKENS is reasonable for 8GB RAM""" # Should be <= 512 for lightweight 3B model self.assertLessEqual(MAX_TOKENS, 512) def test_max_context_turns_set(self): """Test that MAX_CONTEXT_TURNS is reasonable for 8GB RAM""" # Should be <= 8 to control memory self.assertLessEqual(MAX_CONTEXT_TURNS, 8) class TestEpisodeIntegration(unittest.TestCase): """Integration tests for episode execution""" @patch('inference.CloudSOCEnv') @patch('inference.call_llm') def test_episode_basic_flow(self, mock_call_llm, mock_env_class): """Test basic episode flow with mocked environment""" # Mock environment mock_env = MagicMock() mock_env_class.return_value = mock_env mock_env.reset.return_value = ("obs", {"action": "reset"}) mock_env.max_steps = 5 mock_env.get_system_prompt.return_value = "Test system prompt" # Mock steps mock_env.step.side_effect = [ ("obs", 0.1, False, False, {"last_action_error": None}), ("obs", 0.05, True, False, {"last_action_error": None}), # Done ] mock_env.get_state_for_next_task.return_value = None # Mock LLM response valid_action = json.dumps({ "thought": "Test", "tool": "aws.soc.get_alerts", "args": {} }) mock_call_llm.return_value = valid_action # Run episode output = StringIO() with patch('sys.stdout', output): success, steps, rewards, _ = run_episode( task="easy", seed=42, verbose=False ) # Verify output format lines = output.getvalue().strip().split('\n') self.assertTrue(lines[0].startswith("[START]")) self.assertTrue(any(line.startswith("[STEP]") for line in lines)) self.assertTrue(lines[-1].startswith("[END]")) @patch('inference.CloudSOCEnv') @patch('inference.call_llm') def test_episode_handles_parse_errors(self, mock_call_llm, mock_env_class): """Test that episode handles parsing errors gracefully""" mock_env = MagicMock() mock_env_class.return_value = mock_env mock_env.reset.return_value = ("obs", {}) mock_env.max_steps = 3 mock_env.get_system_prompt.return_value = "Test" # First call returns invalid JSON (triggers fallback), second succeeds mock_call_llm.side_effect = [ "completely invalid json", json.dumps({"thought": "Act", "tool": "aws.soc.get_alerts", "args": {}}), ] mock_env.step.side_effect = [ ("obs", -0.02, False, False, {"last_action_error": "PARSE_ERROR"}), ("obs", 0.1, True, False, {"last_action_error": None}), ] mock_env.get_state_for_next_task.return_value = None # Episode should complete without crashing output = StringIO() with patch('sys.stdout', output): success, steps, rewards, _ = run_episode(task="easy", verbose=False) # Should have run 2 steps self.assertEqual(steps, 2) if __name__ == "__main__": import sys # Parse args verbose = "--verbose" in sys.argv if verbose: sys.argv.remove("--verbose") # Run tests suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__]) runner = unittest.TextTestRunner(verbosity=2 if verbose else 1) result = runner.run(suite) sys.exit(0 if result.wasSuccessful() else 1)