move expectations

Signed-off-by: Sammy Oina <[email protected]>
ultravioletrs · Jan 17, 2025 · 83a2b3f · 83a2b3f
1 parent b55a332
commit 83a2b3f
Showing 1 changed file with 8 additions and 12 deletions.
diff --git a/agent/service_test.go b/agent/service_test.go
@@ -37,9 +37,6 @@ const datasetFile = "iris.csv"
 func TestAlgo(t *testing.T) {
 	events := new(mocks.Service)
 
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
 	qp, err := quoteprovider.GetQuoteProvider()
 	require.NoError(t, err)
 
@@ -120,6 +117,9 @@ func TestAlgo(t *testing.T) {
 				metadata.Pairs(algorithm.AlgoTypeKey, tc.algoType, python.PyRuntimeKey, python.PyRuntime),
 			)
 
+			evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			defer evCall.Unset()
+
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
 			svc := New(ctx, mglog.NewMock(), events, qp)
@@ -143,9 +143,6 @@ func TestAlgo(t *testing.T) {
 func TestData(t *testing.T) {
 	events := new(mocks.Service)
 
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
 	qp, err := quoteprovider.GetQuoteProvider()
 	require.NoError(t, err)
 
@@ -212,6 +209,9 @@ func TestData(t *testing.T) {
 					python.PyRuntimeKey, python.PyRuntime),
 			)
 
+			evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			defer evCall.Unset()
+
 			if tc.err != ErrUndeclaredDataset {
 				ctx = IndexToContext(ctx, 0)
 			}
@@ -246,9 +246,6 @@ func TestData(t *testing.T) {
 func TestResult(t *testing.T) {
 	events := new(mocks.Service)
 
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
 	qp, err := quoteprovider.GetQuoteProvider()
 	require.NoError(t, err)
 
@@ -291,6 +288,8 @@ func TestResult(t *testing.T) {
 	}
 
 	for _, tc := range cases {
+		evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+		defer evCall.Unset()
 		t.Run(tc.name, func(t *testing.T) {
 			ctx := metadata.NewIncomingContext(context.Background(),
 				metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
@@ -332,9 +331,6 @@ func TestAttestation(t *testing.T) {
 	events := new(mocks.Service)
 	qp := new(mocks2.QuoteProvider)
 
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
 	cases := []struct {
 		name       string
 		reportData [ReportDataSize]byte