lambdagent/tests/test_dataset.py

"""Tests for lambdagent Dataset."""
from lambdagent.dataset import Dataset


def test_dataset_to_prompt():
    d = Dataset(examples=[("1", "2"), ("2", "3")], description="Successor function")
    prompt = d.to_prompt()
    assert "Successor function" in prompt
    assert "Input: 1" in prompt
    assert "Output: 2" in prompt


def test_dataset_to_lam():
    d = Dataset(examples=[("a", "b")])
    lam = d.to_lam("test_lam")
    assert lam._name == "test_lam"
    assert "Input: a" in lam.prompt