from data.generate_rules import add_rule
from data.generate_scenes import HELD_OUT_SHAPES, generate_scene


def test_scene_generation_is_deterministic():
    first = generate_scene(7, seed=17)
    second = generate_scene(7, seed=17)
    assert first == second
    assert len(first["visual_facts"]) == 5


def test_held_out_rule_changes_combination():
    scene = generate_scene(3, seed=17)
    train = add_rule(scene, seed=17, held_out=False)
    held_out = add_rule(scene, seed=17, held_out=True)
    assert train["combination"] != held_out["combination"]
    assert train["answer"] == held_out["answer"]
    assert len(train["text_facts"]) == 1
    assert len(train["text_rules"]) == 1


def test_ood_scene_uses_held_out_visual_composition():
    scene = generate_scene(5, seed=17, held_out_composition=True)
    first = scene["objects"][0]
    assert first["shape"] == HELD_OUT_SHAPES[first["color"]]