7
7
8
8
"github.com/stretchr/testify/assert"
9
9
"github.com/stretchr/testify/mock"
10
+ "github.com/stretchr/testify/require"
10
11
11
12
"github.com/securego/gosec/v2/issue"
12
13
)
@@ -44,17 +45,16 @@ func TestGenerateSolutionByGemini_Success(t *testing.T) {
44
45
45
46
mockClient := new (MockGenAIClient )
46
47
mockModel := new (MockGenAIGenerativeModel )
47
- mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel )
48
- mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("Autofix for issue 1" , nil )
48
+ mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel ). Once ()
49
+ mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("Autofix for issue 1" , nil ). Once ()
49
50
50
51
// Act
51
52
err := generateSolutionByGemini (mockClient , issues )
52
53
53
54
// Assert
54
- assert .NoError (t , err )
55
- assert .Equal (t , "Autofix for issue 1" , issues [0 ].Autofix )
56
- mockClient .AssertExpectations (t )
57
- mockModel .AssertExpectations (t )
55
+ require .NoError (t , err )
56
+ assert .Equal (t , []* issue.Issue {{What : "Example issue 1" , Autofix : "Autofix for issue 1" }}, issues )
57
+ mock .AssertExpectationsForObjects (t , mockClient , mockModel )
58
58
}
59
59
60
60
func TestGenerateSolutionByGemini_NoCandidates (t * testing.T ) {
@@ -65,17 +65,15 @@ func TestGenerateSolutionByGemini_NoCandidates(t *testing.T) {
65
65
66
66
mockClient := new (MockGenAIClient )
67
67
mockModel := new (MockGenAIGenerativeModel )
68
- mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel )
69
- mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("" , nil )
68
+ mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel ). Once ()
69
+ mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("" , nil ). Once ()
70
70
71
71
// Act
72
72
err := generateSolutionByGemini (mockClient , issues )
73
73
74
74
// Assert
75
- assert .Error (t , err )
76
- assert .Equal (t , "no autofix returned by gemini" , err .Error ())
77
- mockClient .AssertExpectations (t )
78
- mockModel .AssertExpectations (t )
75
+ require .EqualError (t , err , "no autofix returned by gemini" )
76
+ mock .AssertExpectationsForObjects (t , mockClient , mockModel )
79
77
}
80
78
81
79
func TestGenerateSolutionByGemini_APIError (t * testing.T ) {
@@ -86,17 +84,15 @@ func TestGenerateSolutionByGemini_APIError(t *testing.T) {
86
84
87
85
mockClient := new (MockGenAIClient )
88
86
mockModel := new (MockGenAIGenerativeModel )
89
- mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel )
90
- mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("" , errors .New ("API error" ))
87
+ mockClient .On ("GenerativeModel" , GeminiModel ).Return (mockModel ). Once ()
88
+ mockModel .On ("GenerateContent" , mock .Anything , mock .Anything ).Return ("" , errors .New ("API error" )). Once ()
91
89
92
90
// Act
93
91
err := generateSolutionByGemini (mockClient , issues )
94
92
95
93
// Assert
96
- assert .Error (t , err )
97
- assert .Equal (t , "generating autofix with gemini: API error" , err .Error ())
98
- mockClient .AssertExpectations (t )
99
- mockModel .AssertExpectations (t )
94
+ require .EqualError (t , err , "generating autofix with gemini: API error" )
95
+ mock .AssertExpectationsForObjects (t , mockClient , mockModel )
100
96
}
101
97
102
98
func TestGenerateSolution_UnsupportedProvider (t * testing.T ) {
@@ -109,6 +105,5 @@ func TestGenerateSolution_UnsupportedProvider(t *testing.T) {
109
105
err := GenerateSolution ("unsupported-provider" , "test-api-key" , "" , issues )
110
106
111
107
// Assert
112
- assert .Error (t , err )
113
- assert .Equal (t , "ai provider not supported" , err .Error ())
108
+ require .EqualError (t , err , "ai provider not supported" )
114
109
}
0 commit comments