tests/test_env.py

import unittest
import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from environment.env import CodeReviewEnv
from environment.models import ReviewAction, ReviewActionType, Comment, Suggestion


class TestCodeReviewEnv(unittest.TestCase):
    
    def setUp(self):
        self.env = CodeReviewEnv()
    
    def test_reset(self):
        obs = self.env.reset(task_id="bug_detection_easy_1")
        
        self.assertIsNotNone(obs)
        self.assertEqual(obs["task_difficulty"], "easy")
        self.assertEqual(obs["current_step"], 0)
        self.assertIn("code_diff", obs)
        self.assertIn("file_context", obs)
    
    def test_step_add_comment(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.ADD_COMMENT,
            comments=[
                Comment(
                    line_number=3,
                    content="Potential division by zero if list is empty",
                    is_issue=True,
                    severity="high"
                )
            ],
            suggestions=[]
        )
        
        obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertFalse(done)
        self.assertGreaterEqual(reward, 0)
        self.assertTrue(info["last_action_valid"])
        self.assertEqual(len(obs["previous_comments"]), 1)
    
    def test_step_suggest_fix(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.SUGGEST_FIX,
            comments=[],
            suggestions=[
                Suggestion(
                    original_line=3,
                    suggested_code="return total / len(numbers) if numbers else 0",
                    explanation="Add check for empty list"
                )
            ]
        )
        
        obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertFalse(done)
        self.assertEqual(len(obs["previous_suggestions"]), 1)
    
    def test_step_approve(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.APPROVE,
            comments=[],
            suggestions=[],
            final_decision="approved"
        )
        
        obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertTrue(done)
        self.assertEqual(obs["final_decision_made"], "approved")
    
    def test_step_request_changes(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.REQUEST_CHANGES,
            comments=[],
            suggestions=[],
            final_decision="changes_requested"
        )
        
        obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertTrue(done)
        self.assertEqual(obs["final_decision_made"], "changes_requested")
    
    def test_max_steps(self):
        self.env.reset(task_id="bug_detection_easy_1")
        self.env.max_steps = 3
        
        for i in range(3):
            action = ReviewAction(
                action_type=ReviewActionType.ADD_COMMENT,
                comments=[
                    Comment(
                        line_number=3,
                        content=f"Comment {i}",
                        is_issue=True
                    )
                ],
                suggestions=[]
            )
            obs, reward, done, info = self.env.step(action.model_dump())
            
            if i < 2:
                self.assertFalse(done)
            else:
                self.assertTrue(done)
    
    def test_invalid_line_number(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.ADD_COMMENT,
            comments=[
                Comment(
                    line_number=100,
                    content="Invalid line",
                    is_issue=True
                )
            ],
            suggestions=[]
        )
        
        obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertFalse(info["last_action_valid"])
        self.assertIsNotNone(info["error"])
        self.assertEqual(len(obs["previous_comments"]), 0)
    
    def test_get_task_score(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        action = ReviewAction(
            action_type=ReviewActionType.ADD_COMMENT,
            comments=[
                Comment(
                    line_number=3,
                    content="Division by zero issue",
                    is_issue=True,
                    severity="high"
                )
            ],
            suggestions=[
                Suggestion(
                    original_line=3,
                    suggested_code="return total / len(numbers) if numbers else 0",
                    explanation="Fix division by zero"
                )
            ],
            final_decision="changes_requested"
        )
        
        self.env.step(action.model_dump())
        
        score = self.env.get_task_score()
        
        self.assertGreaterEqual(score, 0)
        self.assertLessEqual(score, 1.0)
    
    def test_state_method(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        state = self.env.state()
        
        self.assertIsInstance(state, dict)
        self.assertIn("code_context", state)
        self.assertIn("task_metadata", state)
    
    def test_multiple_actions(self):
        self.env.reset(task_id="bug_detection_easy_1")
        
        actions = [
            ReviewAction(
                action_type=ReviewActionType.ADD_COMMENT,
                comments=[Comment(line_number=3, content="Issue 1", is_issue=True)],
                suggestions=[]
            ),
            ReviewAction(
                action_type=ReviewActionType.SUGGEST_FIX,
                comments=[],
                suggestions=[Suggestion(original_line=3, suggested_code="fix", explanation="Fix")]
            ),
            ReviewAction(
                action_type=ReviewActionType.REQUEST_CHANGES,
                comments=[],
                suggestions=[],
                final_decision="changes_requested"
            )
        ]
        
        for action in actions:
            obs, reward, done, info = self.env.step(action.model_dump())
        
        self.assertTrue(done)
        score = self.env.get_task_score()
        self.assertGreater(score, 0)


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