71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
import pytest
|
|
from src.models import Persona
|
|
|
|
def test_persona_serialization():
|
|
persona = Persona(
|
|
name="SecuritySpecialist",
|
|
provider="anthropic",
|
|
model="claude-3-7-sonnet-20250219",
|
|
preferred_models=["claude-3-7-sonnet-20250219", "claude-3-5-sonnet-20241022"],
|
|
system_prompt="You are a security expert.",
|
|
temperature=0.2,
|
|
top_p=0.9,
|
|
max_output_tokens=4000,
|
|
tool_preset="SecurityTools",
|
|
bias_profile="Execution-Focused"
|
|
)
|
|
|
|
data = persona.to_dict()
|
|
|
|
assert data["provider"] == "anthropic"
|
|
assert data["model"] == "claude-3-7-sonnet-20250219"
|
|
assert "claude-3-5-sonnet-20241022" in data["preferred_models"]
|
|
assert data["system_prompt"] == "You are a security expert."
|
|
assert data["temperature"] == 0.2
|
|
assert data["top_p"] == 0.9
|
|
assert data["max_output_tokens"] == 4000
|
|
assert data["tool_preset"] == "SecurityTools"
|
|
assert data["bias_profile"] == "Execution-Focused"
|
|
|
|
def test_persona_deserialization():
|
|
data = {
|
|
"provider": "gemini",
|
|
"model": "gemini-2.5-flash",
|
|
"preferred_models": ["gemini-2.5-flash"],
|
|
"system_prompt": "You are a helpful assistant.",
|
|
"temperature": 0.5,
|
|
"top_p": 1.0,
|
|
"max_output_tokens": 8192,
|
|
"tool_preset": "Default",
|
|
"bias_profile": "Balanced"
|
|
}
|
|
|
|
persona = Persona.from_dict("Assistant", data)
|
|
|
|
assert persona.name == "Assistant"
|
|
assert persona.provider == "gemini"
|
|
assert persona.model == "gemini-2.5-flash"
|
|
assert persona.preferred_models == ["gemini-2.5-flash"]
|
|
assert persona.system_prompt == "You are a helpful assistant."
|
|
assert persona.temperature == 0.5
|
|
assert persona.top_p == 1.0
|
|
assert persona.max_output_tokens == 8192
|
|
assert persona.tool_preset == "Default"
|
|
assert persona.bias_profile == "Balanced"
|
|
|
|
def test_persona_defaults():
|
|
persona = Persona(name="Minimal", system_prompt="Just the basics")
|
|
assert persona.provider is None
|
|
assert persona.model is None
|
|
assert persona.preferred_models == []
|
|
assert persona.temperature is None
|
|
assert persona.tool_preset is None
|
|
|
|
data = persona.to_dict()
|
|
assert "provider" not in data
|
|
assert "preferred_models" not in data
|
|
assert "temperature" not in data
|
|
|
|
loaded = Persona.from_dict("Minimal", data)
|
|
assert loaded.preferred_models == []
|