77
88 "github.com/stretchr/testify/assert"
99 "github.com/stretchr/testify/mock"
10+ "github.com/stretchr/testify/require"
1011
1112 "github.com/securego/gosec/v2/issue"
1213)
@@ -44,17 +45,16 @@ func TestGenerateSolutionByGemini_Success(t *testing.T) {
4445
4546 mockClient := new (MockGenAIClient )
4647 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 ()
4950
5051 // Act
5152 err := generateSolutionByGemini (mockClient , issues )
5253
5354 // 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 )
5858}
5959
6060func TestGenerateSolutionByGemini_NoCandidates (t * testing.T ) {
@@ -65,17 +65,15 @@ func TestGenerateSolutionByGemini_NoCandidates(t *testing.T) {
6565
6666 mockClient := new (MockGenAIClient )
6767 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 ()
7070
7171 // Act
7272 err := generateSolutionByGemini (mockClient , issues )
7373
7474 // 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 )
7977}
8078
8179func TestGenerateSolutionByGemini_APIError (t * testing.T ) {
@@ -86,17 +84,15 @@ func TestGenerateSolutionByGemini_APIError(t *testing.T) {
8684
8785 mockClient := new (MockGenAIClient )
8886 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 ()
9189
9290 // Act
9391 err := generateSolutionByGemini (mockClient , issues )
9492
9593 // 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 )
10096}
10197
10298func TestGenerateSolution_UnsupportedProvider (t * testing.T ) {
@@ -109,6 +105,5 @@ func TestGenerateSolution_UnsupportedProvider(t *testing.T) {
109105 err := GenerateSolution ("unsupported-provider" , "test-api-key" , "" , issues )
110106
111107 // 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" )
114109}
0 commit comments