test_env.py

"""
Test script for OpenEnv Data Cleaning Environment
"""

from env import DataCleaningEnv
from env.tasks import TaskManager
from env.models import Action, ActionType

# Generate datasets
print("Generating datasets...")
tm = TaskManager()
tm.generate_datasets()
print("Datasets generated successfully")

# Test environment initialization
print("\nTesting environment initialization...")
env = DataCleaningEnv()
task = tm.get_task('easy_001')
obs = env.reset(task)
print('Environment reset successful')
print(f'Total rows: {obs.total_rows}')
print(f'Total columns: {obs.total_columns}')
print(f'Issues remaining: {obs.issues_remaining}')
print(f'Task level: {obs.task_level.value}')
print(f'Task description: {obs.task_description}')

# Test action execution
print("\nTesting action execution...")
action = Action(action_type=ActionType.FILL_MISSING, params={'column': 'age', 'strategy': 'median'})
obs, reward, done, info = env.step(action)
print(f'Action executed: {action.action_type.value}')
print(f'Reward: {reward.total:.4f}')
print(f'Quality: {obs.quality_metrics.overall:.2%}')
print(f'Issues remaining: {obs.issues_remaining}')

# Test another action
action2 = Action(action_type=ActionType.FILL_MISSING, params={'column': 'email', 'strategy': 'mode'})
obs, reward, done, info = env.step(action2)
print(f'\nAction executed: {action2.action_type.value}')
print(f'Reward: {reward.total:.4f}')
print(f'Quality: {obs.quality_metrics.overall:.2%}')
print(f'Issues remaining: {obs.issues_remaining}')

# Test drop duplicates
action3 = Action(action_type=ActionType.DROP_DUPLICATES, params={})
obs, reward, done, info = env.step(action3)
print(f'\nAction executed: {action3.action_type.value}')
print(f'Reward: {reward.total:.4f}')
print(f'Quality: {obs.quality_metrics.overall:.2%}')
print(f'Issues remaining: {obs.issues_remaining}')

# Test submit
action4 = Action(action_type=ActionType.SUBMIT, params={})
obs, reward, done, info = env.step(action4)
print(f'\nAction executed: {action4.action_type.value}')
print(f'Reward: {reward.total:.4f}')
print(f'Quality: {obs.quality_metrics.overall:.2%}')
print(f'Done: {done}')
if 'grade' in info:
    print(f'Final score: {info["grade"].get("final_score", 0):.2%}')

print("\nAll tests passed successfully!")