add seed & size params to RubicsCubeConfig

This commit is contained in:
Andreas Koepf 2025-01-30 00:49:41 +01:00
parent 08ee435d47
commit 25505e3a75
2 changed files with 11 additions and 8 deletions

View file

@ -1,6 +1,5 @@
import pytest
from magiccube.cube import Cube
from reasoning_gym.cognition.rubiks_cube import RubiksCubeConfig, RubiksCubeDataset
@ -17,11 +16,13 @@ def test_rubikscube_config_validation():
def test_rubikscube_deterministic():
"""Test that dataset generates same items with same seed"""
config = RubiksCubeConfig(seed=42, size=15)
config = RubiksCubeConfig(seed=42, size=15) # Only check first 15 entries for speed
dataset1 = RubiksCubeDataset(config)
dataset2 = RubiksCubeDataset(config)
assert len(dataset1) == 15
assert len(dataset2) == 15
for i in range(15): # Only check first 15 entries for speed
for i in range(len(dataset1)):
assert dataset1[i] == dataset2[i]
@ -29,7 +30,8 @@ def test_rubikscube_items():
"""Test basic properties and solution of generated items"""
config = RubiksCubeConfig(
cube_size=3,
scramble_steps=4
scramble_steps=4,
size=100,
)
dataset = RubiksCubeDataset(config)
@ -46,7 +48,6 @@ def test_rubikscube_items():
assert "scramble_moves" in item["metadata"]
assert "example_correct_answer" in item["metadata"]
assert dataset.score_answer(answer=item['metadata']['example_correct_answer'], entry=item) == 1.0
assert dataset.score_answer(answer='R', entry=item) == 0.01
assert dataset.score_answer(answer=item["metadata"]["example_correct_answer"], entry=item) == 1.0
assert dataset.score_answer(answer="R", entry=item) == 0.01
assert dataset.score_answer(answer=None, entry=item) == 0.0