-
Notifications
You must be signed in to change notification settings - Fork 3
/
translate_test.go
133 lines (100 loc) · 3.33 KB
/
translate_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
package dragoman_test
import (
"context"
"testing"
"github.com/MakeNowJust/heredoc/v2"
"github.com/modernice/dragoman"
)
func TestTranslator_Translate(t *testing.T) {
source := heredoc.Docf(`{
"hallo": "Hallo Welt!"
}`)
wantPrompt := heredoc.Docf(`
Translate the following document to English:
---<DOC_BEGIN>---
%s
---<DOC_END>---
Preserve the original document structure and formatting.
Preserve code blocks, placeholders, HTML tags and other structures.
Output only the translated document, no chat.
`, source)
prompt(wantPrompt).expect(t, dragoman.TranslateParams{Document: source})
}
func TestSource(t *testing.T) {
source := heredoc.Docf(`{
"hallo": "Hallo Welt!"
}`)
wantPrompt := heredoc.Docf(`
Translate the following document from French to English:
---<DOC_BEGIN>---
%s
---<DOC_END>---
Preserve the original document structure and formatting.
Preserve code blocks, placeholders, HTML tags and other structures.
Output only the translated document, no chat.
`, source)
prompt(wantPrompt).expect(t, dragoman.TranslateParams{Document: source, Source: "French"})
}
func TestTarget(t *testing.T) {
source := heredoc.Docf(`{
"hallo": "Hallo Welt!"
}`)
wantPrompt := heredoc.Docf(`
Translate the following document to French:
---<DOC_BEGIN>---
%s
---<DOC_END>---
Preserve the original document structure and formatting.
Preserve code blocks, placeholders, HTML tags and other structures.
Output only the translated document, no chat.
`, source)
prompt(wantPrompt).expect(t, dragoman.TranslateParams{Document: source, Target: "French"})
}
func TestPreserve(t *testing.T) {
source := heredoc.Docf(`{
"hallo": "Hallo, ich bin der HalloWeltBot!"
}`)
wantPrompt := heredoc.Docf(`
Translate the following document to English:
---<DOC_BEGIN>---
%s
---<DOC_END>---
Preserve the original document structure and formatting.
Preserve code blocks, placeholders, HTML tags and other structures.
Do not translate the following terms: HalloWeltBot
Output only the translated document, no chat.
`, source)
prompt(wantPrompt).expect(t, dragoman.TranslateParams{Document: source, Preserve: []string{"HalloWeltBot"}})
}
func TestPreserve_multiple(t *testing.T) {
source := heredoc.Docf(`{
"hallo": "Hallo, ich bin der HalloWeltBot aus der WeltFabrik!"
}`)
wantPrompt := heredoc.Docf(`
Translate the following document to English:
---<DOC_BEGIN>---
%s
---<DOC_END>---
Preserve the original document structure and formatting.
Preserve code blocks, placeholders, HTML tags and other structures.
Do not translate the following terms: HalloWeltBot, WeltFabrik
Output only the translated document, no chat.
`, source)
prompt(wantPrompt).expect(t, dragoman.TranslateParams{Document: source, Preserve: []string{"HalloWeltBot", "WeltFabrik"}})
}
type prompt string
func (p prompt) expect(t *testing.T, params dragoman.TranslateParams) {
t.Helper()
var providedPrompt string
model := dragoman.ModelFunc(func(_ context.Context, prompt string) (string, error) {
providedPrompt = prompt
return "", nil
})
trans := dragoman.NewTranslator(model)
if _, err := trans.Translate(context.Background(), params); err != nil {
t.Fatalf("unexpected error: %v", err)
}
if providedPrompt != string(p) {
t.Errorf("expected prompt to be\n\n%s\n\nbut prompt was\n\n%s", p, providedPrompt)
}
}