Skip to content

Commit 9ea0b0b

Browse files
dgageotdgageot
authored
Merge pull request #315 from dgageot/test-assertions
Use assertions in tests
2 parents cf92f8d + a23c0d3 commit 9ea0b0b

11 files changed

Lines changed: 220 additions & 566 deletions

File tree

‎pkg/content/store_test.go‎

Lines changed: 24 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -8,54 +8,36 @@ import (
88
"github.com/google/go-containerregistry/pkg/v1/mutate"
99
"github.com/google/go-containerregistry/pkg/v1/static"
1010
"github.com/google/go-containerregistry/pkg/v1/types"
11+
"github.com/stretchr/testify/assert"
12+
"github.com/stretchr/testify/require"
1113
)
1214

1315
func TestStoreBasicOperations(t *testing.T) {
1416
store, err := NewStore(WithBaseDir(t.TempDir()))
15-
if err != nil {
16-
t.Fatalf("Failed to create store: %v", err)
17-
}
17+
require.NoError(t, err)
1818

1919
testData := []byte("Hello, World! This is a test artifact.")
2020
layer := static.NewLayer(testData, types.OCIUncompressedLayer)
2121
img := empty.Image
2222
img, err = mutate.AppendLayers(img, layer)
23-
if err != nil {
24-
t.Fatalf("Failed to create test image: %v", err)
25-
}
23+
require.NoError(t, err)
2624

2725
testRef := "hello-world:v1.0.0"
2826
digest, err := store.StoreArtifact(img, testRef)
29-
if err != nil {
30-
t.Fatalf("Failed to store artifact: %v", err)
31-
}
27+
require.NoError(t, err)
3228

3329
retrievedImg, err := store.GetArtifactImage(testRef)
34-
if err != nil {
35-
t.Fatalf("Failed to retrieve artifact by reference: %v", err)
36-
}
37-
38-
if retrievedImg == nil {
39-
t.Fatal("Retrieved image is nil")
40-
}
30+
require.NoError(t, err)
31+
assert.NotNil(t, retrievedImg)
4132

4233
metadata, err := store.GetArtifactMetadata(testRef)
43-
if err != nil {
44-
t.Fatalf("Failed to get metadata: %v", err)
45-
}
46-
47-
if metadata.Reference != testRef {
48-
t.Errorf("Expected reference %s, got %s", testRef, metadata.Reference)
49-
}
34+
require.NoError(t, err)
5035

51-
if metadata.Digest != digest {
52-
t.Errorf("Expected digest %s, got %s", digest, metadata.Digest)
53-
}
36+
assert.Equal(t, testRef, metadata.Reference)
37+
assert.Equal(t, digest, metadata.Digest)
5438

5539
artifacts, err := store.ListArtifacts()
56-
if err != nil {
57-
t.Fatalf("Failed to list artifacts: %v", err)
58-
}
40+
require.NoError(t, err)
5941

6042
found := false
6143
for _, artifact := range artifacts {
@@ -65,16 +47,12 @@ func TestStoreBasicOperations(t *testing.T) {
6547
}
6648
}
6749

68-
if !found {
69-
t.Error("Artifact not found in list")
70-
}
50+
assert.True(t, found, "Artifact not found in list")
7151
}
7252

7353
func TestStoreMultipleArtifacts(t *testing.T) {
7454
store, err := NewStore(WithBaseDir(t.TempDir()))
75-
if err != nil {
76-
t.Fatalf("Failed to create store: %v", err)
77-
}
55+
require.NoError(t, err)
7856

7957
testRefs := []string{
8058
"app1:v1.0.0",
@@ -89,69 +67,47 @@ func TestStoreMultipleArtifacts(t *testing.T) {
8967
layer := static.NewLayer(testData, types.OCIUncompressedLayer)
9068
img := empty.Image
9169
img, err = mutate.AppendLayers(img, layer)
92-
if err != nil {
93-
t.Fatalf("Failed to create test image %d: %v", i+1, err)
94-
}
70+
require.NoError(t, err)
9571

9672
digest, err := store.StoreArtifact(img, ref)
97-
if err != nil {
98-
t.Fatalf("Failed to store artifact %d: %v", i+1, err)
99-
}
73+
require.NoError(t, err)
10074

10175
digests = append(digests, digest)
10276
}
10377

10478
artifacts, err := store.ListArtifacts()
105-
if err != nil {
106-
t.Fatalf("Failed to list artifacts: %v", err)
107-
}
79+
require.NoError(t, err)
10880

109-
if len(artifacts) < len(testRefs) {
110-
t.Errorf("Expected at least %d artifacts, got %d", len(testRefs), len(artifacts))
111-
}
81+
assert.GreaterOrEqual(t, len(artifacts), len(testRefs))
11282

11383
for _, ref := range testRefs {
11484
img, err := store.GetArtifactImage(ref)
115-
if err != nil {
116-
t.Errorf("Failed to retrieve artifact %s: %v", ref, err)
117-
}
118-
if img == nil {
119-
t.Errorf("Retrieved image for %s is nil", ref)
120-
}
85+
require.NoError(t, err)
86+
assert.NotNil(t, img)
12187
}
12288
}
12389

12490
func TestStoreResolution(t *testing.T) {
12591
store, err := NewStore(WithBaseDir(t.TempDir()))
126-
if err != nil {
127-
t.Fatalf("Failed to create store: %v", err)
128-
}
92+
require.NoError(t, err)
12993

13094
testData := []byte("Resolution test artifact")
13195
layer := static.NewLayer(testData, types.OCIUncompressedLayer)
13296
img := empty.Image
13397
img, err = mutate.AppendLayers(img, layer)
134-
if err != nil {
135-
t.Fatalf("Failed to create test image: %v", err)
136-
}
98+
require.NoError(t, err)
13799

138100
testRef := "resolution-test:latest"
139101
_, err = store.StoreArtifact(img, testRef)
140-
if err != nil {
141-
t.Fatalf("Failed to store artifact: %v", err)
142-
}
102+
require.NoError(t, err)
143103

144104
testCases := []string{
145105
testRef,
146106
}
147107

148108
for _, tc := range testCases {
149109
img, err := store.GetArtifactImage(tc)
150-
if err != nil {
151-
t.Errorf("Failed to resolve identifier %s: %v", tc, err)
152-
}
153-
if img == nil {
154-
t.Errorf("Retrieved image for %s is nil", tc)
155-
}
110+
require.NoError(t, err)
111+
assert.NotNil(t, img)
156112
}
157113
}

‎pkg/model/provider/anthropic/client_test.go‎

Lines changed: 42 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ import (
55
"strings"
66
"testing"
77

8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
811
"github.com/docker/cagent/pkg/chat"
912
"github.com/docker/cagent/pkg/tools"
1013
)
@@ -16,9 +19,7 @@ func TestConvertMessages_SkipEmptySystemText(t *testing.T) {
1619
}}
1720

1821
out := convertMessages(msgs)
19-
if len(out) != 0 {
20-
t.Fatalf("expected 0 messages, got %d", len(out))
21-
}
22+
assert.Empty(t, out)
2223
}
2324

2425
func TestConvertMessages_SkipEmptyUserText_NoMultiContent(t *testing.T) {
@@ -28,9 +29,7 @@ func TestConvertMessages_SkipEmptyUserText_NoMultiContent(t *testing.T) {
2829
}}
2930

3031
out := convertMessages(msgs)
31-
if len(out) != 0 {
32-
t.Fatalf("expected 0 messages, got %d", len(out))
33-
}
32+
assert.Empty(t, out)
3433
}
3534

3635
func TestConvertMessages_UserMultiContent_SkipEmptyText_KeepImage(t *testing.T) {
@@ -43,34 +42,23 @@ func TestConvertMessages_UserMultiContent_SkipEmptyText_KeepImage(t *testing.T)
4342
}}
4443

4544
out := convertMessages(msgs)
46-
if len(out) != 1 {
47-
t.Fatalf("expected 1 message, got %d", len(out))
48-
}
45+
require.Len(t, out, 1)
4946

5047
b, err := json.Marshal(out[0])
51-
if err != nil {
52-
t.Fatalf("marshal error: %v", err)
53-
}
48+
require.NoError(t, err)
5449
// Basic JSON structure checks
5550
var m map[string]any
56-
if err := json.Unmarshal(b, &m); err != nil {
57-
t.Fatalf("unmarshal error: %v", err)
58-
}
51+
require.NoError(t, json.Unmarshal(b, &m))
5952
// role should be user
60-
if role, _ := m["role"].(string); role != "user" {
61-
t.Fatalf("expected role 'user', got %v", m["role"])
62-
}
53+
assert.Equal(t, "user", m["role"])
6354
// content should contain exactly one block (the image)
64-
if content, _ := m["content"].([]any); len(content) != 1 {
65-
t.Fatalf("expected 1 content block, got %d", len(content))
66-
}
55+
content, ok := m["content"].([]any)
56+
require.True(t, ok)
57+
assert.Len(t, content, 1)
6758
// and it should be an image block
68-
if content, _ := m["content"].([]any); len(content) == 1 {
69-
cb, _ := content[0].(map[string]any)
70-
if typ, _ := cb["type"].(string); typ != "image" {
71-
t.Fatalf("expected content block type 'image', got %v", typ)
72-
}
73-
}
59+
cb, ok := content[0].(map[string]any)
60+
require.True(t, ok)
61+
assert.Equal(t, "image", cb["type"])
7462
}
7563

7664
func TestConvertMessages_SkipEmptyAssistantText_NoToolCalls(t *testing.T) {
@@ -80,9 +68,7 @@ func TestConvertMessages_SkipEmptyAssistantText_NoToolCalls(t *testing.T) {
8068
}}
8169

8270
out := convertMessages(msgs)
83-
if len(out) != 0 {
84-
t.Fatalf("expected 0 messages, got %d", len(out))
85-
}
71+
assert.Empty(t, out)
8672
}
8773

8874
func TestConvertMessages_AssistantToolCalls_NoText_IncludesToolUse(t *testing.T) {
@@ -95,29 +81,19 @@ func TestConvertMessages_AssistantToolCalls_NoText_IncludesToolUse(t *testing.T)
9581
}}
9682

9783
out := convertMessages(msgs)
98-
if len(out) != 1 {
99-
t.Fatalf("expected 1 message, got %d", len(out))
100-
}
84+
require.Len(t, out, 1)
10185

10286
b, err := json.Marshal(out[0])
103-
if err != nil {
104-
t.Fatalf("marshal error: %v", err)
105-
}
87+
require.NoError(t, err)
10688
var m map[string]any
107-
if err := json.Unmarshal(b, &m); err != nil {
108-
t.Fatalf("unmarshal error: %v", err)
109-
}
110-
if role, _ := m["role"].(string); role != "assistant" {
111-
t.Fatalf("expected role 'assistant', got %v", m["role"])
112-
}
113-
content, _ := m["content"].([]any)
114-
if len(content) != 1 {
115-
t.Fatalf("expected 1 content block, got %d", len(content))
116-
}
117-
cb, _ := content[0].(map[string]any)
118-
if typ, _ := cb["type"].(string); typ != "tool_use" {
119-
t.Fatalf("expected content block type 'tool_use', got %v", typ)
120-
}
89+
require.NoError(t, json.Unmarshal(b, &m))
90+
assert.Equal(t, "assistant", m["role"])
91+
content, ok := m["content"].([]any)
92+
require.True(t, ok)
93+
assert.Len(t, content, 1)
94+
cb, ok := content[0].(map[string]any)
95+
require.True(t, ok)
96+
assert.Equal(t, "tool_use", cb["type"])
12197
}
12298

12399
func TestSystemMessages_AreExtractedAndNotInMessageList(t *testing.T) {
@@ -128,18 +104,12 @@ func TestSystemMessages_AreExtractedAndNotInMessageList(t *testing.T) {
128104

129105
// System blocks should be extracted
130106
sys := extractSystemBlocks(msgs)
131-
if len(sys) != 1 {
132-
t.Fatalf("expected 1 system block, got %d", len(sys))
133-
}
134-
if strings.TrimSpace(sys[0].Text) != "system rules here" {
135-
t.Fatalf("unexpected system text: %q", sys[0].Text)
136-
}
107+
require.Len(t, sys, 1)
108+
assert.Equal(t, "system rules here", strings.TrimSpace(sys[0].Text))
137109

138110
// System role messages must not appear in the anthropic messages list
139111
out := convertMessages(msgs)
140-
if len(out) != 1 {
141-
t.Fatalf("expected 1 non-system message, got %d", len(out))
142-
}
112+
assert.Len(t, out, 1)
143113
}
144114

145115
func TestSystemMessages_MultipleExtractedAndExcludedFromMessageList(t *testing.T) {
@@ -150,20 +120,12 @@ func TestSystemMessages_MultipleExtractedAndExcludedFromMessageList(t *testing.T
150120
}
151121

152122
sys := extractSystemBlocks(msgs)
153-
if len(sys) != 2 {
154-
t.Fatalf("expected 2 system blocks, got %d", len(sys))
155-
}
156-
if strings.TrimSpace(sys[0].Text) != "sys A" {
157-
t.Fatalf("unexpected first system text: %q", sys[0].Text)
158-
}
159-
if strings.TrimSpace(sys[1].Text) != "sys B" {
160-
t.Fatalf("unexpected second system text: %q", sys[1].Text)
161-
}
123+
require.Len(t, sys, 2)
124+
assert.Equal(t, "sys A", strings.TrimSpace(sys[0].Text))
125+
assert.Equal(t, "sys B", strings.TrimSpace(sys[1].Text))
162126

163127
out := convertMessages(msgs)
164-
if len(out) != 1 {
165-
t.Fatalf("expected 1 non-system message, got %d", len(out))
166-
}
128+
assert.Len(t, out, 1)
167129
}
168130

169131
func TestSystemMessages_InterspersedExtractedAndExcluded(t *testing.T) {
@@ -177,33 +139,20 @@ func TestSystemMessages_InterspersedExtractedAndExcluded(t *testing.T) {
177139

178140
// All system messages should be extracted in order of appearance
179141
sys := extractSystemBlocks(msgs)
180-
if len(sys) != 2 {
181-
t.Fatalf("expected 2 system blocks, got %d", len(sys))
182-
}
183-
if strings.TrimSpace(sys[0].Text) != "S1" {
184-
t.Fatalf("unexpected first system text: %q", sys[0].Text)
185-
}
186-
if strings.TrimSpace(sys[1].Text) != "S2" {
187-
t.Fatalf("unexpected second system text: %q", sys[1].Text)
188-
}
142+
require.Len(t, sys, 2)
143+
assert.Equal(t, "S1", strings.TrimSpace(sys[0].Text))
144+
assert.Equal(t, "S2", strings.TrimSpace(sys[1].Text))
189145

190146
// Converted messages must exclude system roles and preserve order of others
191147
out := convertMessages(msgs)
192-
if len(out) != 3 {
193-
t.Fatalf("expected 3 non-system messages, got %d", len(out))
194-
}
148+
require.Len(t, out, 3)
195149
// Check roles: user, assistant, user
196-
for i, expected := range []string{"user", "assistant", "user"} {
150+
expectedRoles := []string{"user", "assistant", "user"}
151+
for i, expected := range expectedRoles {
197152
b, err := json.Marshal(out[i])
198-
if err != nil {
199-
t.Fatalf("marshal error: %v", err)
200-
}
153+
require.NoError(t, err)
201154
var m map[string]any
202-
if err := json.Unmarshal(b, &m); err != nil {
203-
t.Fatalf("unmarshal error: %v", err)
204-
}
205-
if role, _ := m["role"].(string); role != expected {
206-
t.Fatalf("unexpected role at %d: got %q want %q", i, role, expected)
207-
}
155+
require.NoError(t, json.Unmarshal(b, &m))
156+
assert.Equal(t, expected, m["role"])
208157
}
209158
}

0 commit comments

Comments
 (0)