|
1
|
"""Tests for prompt template management.""" |
|
2
|
|
|
3
|
from video_processor.utils.prompt_templates import ( |
|
4
|
DEFAULT_TEMPLATES, |
|
5
|
PromptTemplate, |
|
6
|
default_prompt_manager, |
|
7
|
) |
|
8
|
|
|
9
|
|
|
10
|
class TestPromptTemplate: |
|
11
|
def test_default_templates_loaded(self): |
|
12
|
pm = PromptTemplate(default_templates=DEFAULT_TEMPLATES) |
|
13
|
assert len(pm.templates) == 10 |
|
14
|
|
|
15
|
def test_all_expected_templates_exist(self): |
|
16
|
expected = [ |
|
17
|
"content_analysis", |
|
18
|
"diagram_extraction", |
|
19
|
"action_item_detection", |
|
20
|
"content_summary", |
|
21
|
"summary_generation", |
|
22
|
"key_points_extraction", |
|
23
|
"entity_extraction", |
|
24
|
"relationship_extraction", |
|
25
|
"diagram_analysis", |
|
26
|
"mermaid_generation", |
|
27
|
] |
|
28
|
for name in expected: |
|
29
|
assert name in DEFAULT_TEMPLATES, f"Missing template: {name}" |
|
30
|
|
|
31
|
def test_get_template(self): |
|
32
|
pm = PromptTemplate(default_templates={"test": "Hello $name"}) |
|
33
|
template = pm.get_template("test") |
|
34
|
assert template is not None |
|
35
|
|
|
36
|
def test_get_missing_template(self): |
|
37
|
pm = PromptTemplate(default_templates={}) |
|
38
|
assert pm.get_template("nonexistent") is None |
|
39
|
|
|
40
|
def test_format_prompt(self): |
|
41
|
pm = PromptTemplate(default_templates={"greet": "Hello $name, welcome to $place"}) |
|
42
|
result = pm.format_prompt("greet", name="Alice", place="Wonderland") |
|
43
|
assert "Alice" in result |
|
44
|
assert "Wonderland" in result |
|
45
|
|
|
46
|
def test_format_missing_template(self): |
|
47
|
pm = PromptTemplate(default_templates={}) |
|
48
|
result = pm.format_prompt("nonexistent", key="value") |
|
49
|
assert result is None |
|
50
|
|
|
51
|
def test_safe_substitute_missing_vars(self): |
|
52
|
pm = PromptTemplate(default_templates={"test": "Hello $name and $other"}) |
|
53
|
result = pm.format_prompt("test", name="Alice") |
|
54
|
assert "Alice" in result |
|
55
|
assert "$other" in result # safe_substitute keeps unresolved vars |
|
56
|
|
|
57
|
def test_add_template(self): |
|
58
|
pm = PromptTemplate(default_templates={}) |
|
59
|
pm.add_template("new", "New template: $var") |
|
60
|
result = pm.format_prompt("new", var="value") |
|
61
|
assert "value" in result |
|
62
|
|
|
63
|
def test_save_template_no_dir(self): |
|
64
|
pm = PromptTemplate(default_templates={"test": "content"}) |
|
65
|
assert pm.save_template("test") is False |
|
66
|
|
|
67
|
def test_save_template_missing_name(self): |
|
68
|
pm = PromptTemplate(default_templates={}) |
|
69
|
assert pm.save_template("nonexistent") is False |
|
70
|
|
|
71
|
def test_save_and_load_from_dir(self, tmp_path): |
|
72
|
templates_dir = tmp_path / "templates" |
|
73
|
templates_dir.mkdir() |
|
74
|
(templates_dir / "custom.txt").write_text("Custom: $data") |
|
75
|
|
|
76
|
pm = PromptTemplate(templates_dir=templates_dir) |
|
77
|
assert "custom" in pm.templates |
|
78
|
result = pm.format_prompt("custom", data="hello") |
|
79
|
assert "hello" in result |
|
80
|
|
|
81
|
def test_save_template_to_dir(self, tmp_path): |
|
82
|
templates_dir = tmp_path / "templates" |
|
83
|
pm = PromptTemplate( |
|
84
|
templates_dir=templates_dir, |
|
85
|
default_templates={"saveme": "Save this: $x"}, |
|
86
|
) |
|
87
|
result = pm.save_template("saveme") |
|
88
|
assert result is True |
|
89
|
assert (templates_dir / "saveme.txt").exists() |
|
90
|
|
|
91
|
|
|
92
|
class TestDefaultPromptManager: |
|
93
|
def test_is_initialized(self): |
|
94
|
assert default_prompt_manager is not None |
|
95
|
assert len(default_prompt_manager.templates) == 10 |
|
96
|
|
|
97
|
def test_entity_extraction_template_has_content_var(self): |
|
98
|
result = default_prompt_manager.format_prompt( |
|
99
|
"entity_extraction", content="some transcript" |
|
100
|
) |
|
101
|
assert "some transcript" in result |
|
102
|
|
|
103
|
def test_mermaid_generation_template(self): |
|
104
|
result = default_prompt_manager.format_prompt( |
|
105
|
"mermaid_generation", |
|
106
|
diagram_type="flowchart", |
|
107
|
text_content="A -> B", |
|
108
|
semantic_analysis="Flow diagram", |
|
109
|
) |
|
110
|
assert "flowchart" in result |
|
111
|
assert "A -> B" in result |
|
112
|
|