From 1cd2a15f52c89d128834485aa53fa64ace6fa286 Mon Sep 17 00:00:00 2001
From: Software Developer <7852635+dsuhinin@users.noreply.github.com>
Date: Mon, 11 Dec 2023 18:41:26 +0100
Subject: [PATCH] Improve unit tests according to changes in
 DATA-DOG/go-sqlmock. (#729)

Improve unit tests according to changes in DATA-DOG/go-sqlmock.
---
 pkg/api/mlflow/dao/repositories/run_test.go | 24 ++++++++++++---------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/pkg/api/mlflow/dao/repositories/run_test.go b/pkg/api/mlflow/dao/repositories/run_test.go
index 90878e9c5..9a18569df 100644
--- a/pkg/api/mlflow/dao/repositories/run_test.go
+++ b/pkg/api/mlflow/dao/repositories/run_test.go
@@ -38,11 +38,6 @@ func Test_renumberRows(t *testing.T) {
 	//nolint:errcheck
 	defer mockDb.Close()
 
-	lockExpect := func() {
-		mock.ExpectExec("LOCK TABLE runs").WillReturnResult(sqlmock.NewResult(0, 1))
-		mock.ExpectExec(`UPDATE runs`).WillReturnResult(sqlmock.NewResult(0, 1))
-	}
-
 	dialector := postgres.New(postgres.Config{
 		Conn:       mockDb,
 		DriverName: "postgres",
@@ -55,12 +50,21 @@ func Test_renumberRows(t *testing.T) {
 	for _, tc := range testData {
 		t.Run(tc.name, func(t *testing.T) {
 			if tc.startWith < 0 {
-				err := repo.renumberRows(db, tc.startWith)
-				assert.EqualError(t, err, "attempting to renumber with less than 0 row number value")
+				assert.EqualError(
+					t,
+					repo.renumberRows(db, tc.startWith),
+					"attempting to renumber with less than 0 row number value",
+				)
 			} else {
-				lockExpect()
-				err := repo.renumberRows(db, tc.startWith)
-				assert.NoError(t, err)
+				mock.ExpectExec(
+					"LOCK TABLE runs",
+				).WillReturnResult(sqlmock.NewResult(0, 1))
+				mock.ExpectExec(`UPDATE runs`).WithArgs(
+					tc.startWith, tc.startWith,
+				).WillReturnResult(
+					sqlmock.NewResult(0, 1),
+				)
+				assert.NoError(t, repo.renumberRows(db, tc.startWith))
 			}
 		})
 	}