muffin2006/document-classification-env

1"""2Test script to validate the Document Classification Environment3"""4 5import sys6from environment import DocumentClassificationEnv7from grading import BaselineAgent, AgentGrader8 9 10def test_environment_creation():11    """Test that environments can be created for all difficulties"""12    print("Testing environment creation...")13    for difficulty in ["easy", "medium", "hard"]:14        try:15            env = DocumentClassificationEnv(task_difficulty=difficulty)16            obs, info = env.reset()17            print(f"✓ {difficulty} environment created successfully")18            assert "features" in obs, f"Missing features in {difficulty} observation"19            assert obs["features"].shape[0] > 0, f"Invalid feature shape in {difficulty}"20        except Exception as e:21            print(f"✗ Failed to create {difficulty} environment: {e}")22            return False23    return True24 25 26def test_step_function():27    """Test that step function works correctly"""28    print("\nTesting step function...")29    env = DocumentClassificationEnv(task_difficulty="easy")30    obs, _ = env.reset()31    32    for i in range(5):33        action = env.action_space.sample()34        obs, reward, done, truncated, info = env.step(action)35        assert isinstance(reward, float), f"Reward should be float, got {type(reward)}"36        assert isinstance(done, bool), f"Done should be bool, got {type(done)}"37        print(f"✓ Step {i+1}: reward={reward:.3f}, accuracy={info['episode_accuracy']:.3f}")38        if done:39            break40    41    print(f"✓ Step function working correctly")42    return True43 44 45def test_state_function():46    """Test that state function works correctly"""47    print("\nTesting state function...")48    env = DocumentClassificationEnv(task_difficulty="easy")49    obs, _ = env.reset()50    51    for _ in range(3):52        action = env.action_space.sample()53        obs, reward, done, truncated, info = env.step(action)54    55    state = env.state()56    assert "current_observation" in state, "Missing current_observation in state"57    assert "current_document_index" in state, "Missing current_document_index in state"58    assert "episode_reward_total" in state, "Missing episode_reward_total in state"59    print(f"✓ State function working correctly")60    print(f"  State keys: {list(state.keys())}")61    return True62 63 64def test_baseline_agent():65    """Test baseline agent on each difficulty"""66    print("\nTesting baseline agent...")67    for difficulty in ["easy"]:  # Test only easy for speed68        try:69            env = DocumentClassificationEnv(task_difficulty=difficulty)70            agent = BaselineAgent(difficulty)71            obs, _ = env.reset()72            73            for _ in range(5):74                action = agent.decide(obs)75                assert 0 <= action < env.num_categories, f"Invalid action: {action}"76                obs, reward, done, truncated, info = env.step(action)77                if done:78                    break79            80            print(f"✓ Baseline agent working on {difficulty}")81        except Exception as e:82            print(f"✗ Baseline agent failed on {difficulty}: {e}")83            return False84    return True85 86 87def test_grading():88    """Test the grading system"""89    print("\nTesting grading system...")90    try:91        grader = AgentGrader("easy")92        agent = BaselineAgent("easy")93        94        score, metrics = grader.grade_agent(agent.decide, verbose=False)95        96        assert 0.0 <= score <= 1.0, f"Score out of range: {score}"97        assert "accuracy" in metrics, "Missing accuracy in metrics"98        assert metrics["accuracy"] >= 0.0, "Negative accuracy"99        100        print(f"✓ Grading system working")101        print(f"  Score: {score:.4f}")102        print(f"  Accuracy: {metrics['accuracy']:.4f}")103        return True104    except Exception as e:105        print(f"✗ Grading system failed: {e}")106        return False107 108 109def run_all_tests():110    """Run all tests"""111    print("="*60)112    print("Document Classification Environment - Test Suite")113    print("="*60)114    115    tests = [116        ("Environment Creation", test_environment_creation),117        ("Step Function", test_step_function),118        ("State Function", test_state_function),119        ("Baseline Agent", test_baseline_agent),120        ("Grading System", test_grading),121    ]122    123    results = []124    for test_name, test_func in tests:125        try:126            result = test_func()127            results.append((test_name, result))128        except Exception as e:129            print(f"✗ {test_name} failed with exception: {e}")130            results.append((test_name, False))131    132    # Summary133    print("\n" + "="*60)134    print("TEST SUMMARY")135    print("="*60)136    137    passed = sum(1 for _, result in results if result)138    total = len(results)139    140    for test_name, result in results:141        status = "✓ PASS" if result else "✗ FAIL"142        print(f"{status} - {test_name}")143    144    print(f"\nTotal: {passed}/{total} tests passed")145    print("="*60)146    147    return passed == total148 149 150if __name__ == "__main__":151    success = run_all_tests()152    sys.exit(0 if success else 1)153