60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
import pytest
|
|
from src.synthesis_formatter import format_takes_diff
|
|
|
|
def test_format_takes_diff_empty():
|
|
assert format_takes_diff({}) == ""
|
|
|
|
def test_format_takes_diff_single_take():
|
|
takes = {
|
|
"take1": [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi"}
|
|
]
|
|
}
|
|
expected = "=== Shared History ===\nuser: hello\nassistant: hi\n\n=== Variations ===\n"
|
|
assert format_takes_diff(takes) == expected
|
|
|
|
def test_format_takes_diff_common_prefix():
|
|
takes = {
|
|
"take1": [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi"},
|
|
{"role": "user", "content": "how are you?"},
|
|
{"role": "assistant", "content": "I am fine."}
|
|
],
|
|
"take2": [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi"},
|
|
{"role": "user", "content": "what is the time?"},
|
|
{"role": "assistant", "content": "It is noon."}
|
|
]
|
|
}
|
|
expected = (
|
|
"=== Shared History ===\n"
|
|
"user: hello\n"
|
|
"assistant: hi\n\n"
|
|
"=== Variations ===\n"
|
|
"[take1]\n"
|
|
"user: how are you?\n"
|
|
"assistant: I am fine.\n\n"
|
|
"[take2]\n"
|
|
"user: what is the time?\n"
|
|
"assistant: It is noon.\n"
|
|
)
|
|
assert format_takes_diff(takes) == expected
|
|
|
|
def test_format_takes_diff_no_common_prefix():
|
|
takes = {
|
|
"take1": [{"role": "user", "content": "a"}],
|
|
"take2": [{"role": "user", "content": "b"}]
|
|
}
|
|
expected = (
|
|
"=== Shared History ===\n\n"
|
|
"=== Variations ===\n"
|
|
"[take1]\n"
|
|
"user: a\n\n"
|
|
"[take2]\n"
|
|
"user: b\n"
|
|
)
|
|
assert format_takes_diff(takes) == expected
|