chore: refactor grpc tests to use a better mocking api

Author: Davidonium

internal/cli/config_load_test.go

@@ -1,7 +1,6 @@
 package cli_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -16,51 +15,46 @@ import (
 // loaded via --config reaches the gRPC request.
 func TestConfigLoad_FlagSetsAccountID(t *testing.T) {
 	env := testutil.NewBareTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
 	cfgPath := testutil.WriteConfigFile(t, t.TempDir(), map[string]any{
 		"account_id": "account-from-file",
 	})
 
-	var capturedAccountID string
-	env.Server.ListClustersFunc = func(_ context.Context, req *clusterv1.ListClustersRequest) (*clusterv1.ListClustersResponse, error) {
-		capturedAccountID = req.GetAccountId()
-		return &clusterv1.ListClustersResponse{}, nil
-	}
+	env.Server.ListClustersCalls.Returns(&clusterv1.ListClustersResponse{}, nil)
 
 	_, _, err := testutil.Exec(t, env, "--config", cfgPath, "cluster", "list")
 	require.NoError(t, err)
-	assert.Equal(t, "account-from-file", capturedAccountID)
+
+	req, ok := env.Server.ListClustersCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "account-from-file", req.GetAccountId())
 }
 
 // TestConfigLoad_EnvVarSetsAccountID verifies that QDRANT_CLOUD_CONFIG env var
 // is respected when no --config flag is given.
 func TestConfigLoad_EnvVarSetsAccountID(t *testing.T) {
 	env := testutil.NewBareTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
 	cfgPath := testutil.WriteConfigFile(t, t.TempDir(), map[string]any{
 		"account_id": "account-from-envvar",
 	})
 
 	t.Setenv("QDRANT_CLOUD_CONFIG", cfgPath)
 
-	var capturedAccountID string
-	env.Server.ListClustersFunc = func(_ context.Context, req *clusterv1.ListClustersRequest) (*clusterv1.ListClustersResponse, error) {
-		capturedAccountID = req.GetAccountId()
-		return &clusterv1.ListClustersResponse{}, nil
-	}
+	env.Server.ListClustersCalls.Returns(&clusterv1.ListClustersResponse{}, nil)
 
 	_, _, err := testutil.Exec(t, env, "cluster", "list")
 	require.NoError(t, err)
-	assert.Equal(t, "account-from-envvar", capturedAccountID)
+
+	req, ok := env.Server.ListClustersCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "account-from-envvar", req.GetAccountId())
 }
 
 // TestConfigLoad_FlagOverridesEnvVar verifies that --config flag takes
 // precedence over QDRANT_CLOUD_CONFIG env var.
 func TestConfigLoad_FlagOverridesEnvVar(t *testing.T) {
 	env := testutil.NewBareTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
 	dir := t.TempDir()
 	flagCfg := testutil.WriteConfigFile(t, dir, map[string]any{
@@ -73,34 +67,31 @@ func TestConfigLoad_FlagOverridesEnvVar(t *testing.T) {
 
 	t.Setenv("QDRANT_CLOUD_CONFIG", envCfg)
 
-	var capturedAccountID string
-	env.Server.ListClustersFunc = func(_ context.Context, req *clusterv1.ListClustersRequest) (*clusterv1.ListClustersResponse, error) {
-		capturedAccountID = req.GetAccountId()
-		return &clusterv1.ListClustersResponse{}, nil
-	}
+	env.Server.ListClustersCalls.Returns(&clusterv1.ListClustersResponse{}, nil)
 
 	_, _, err := testutil.Exec(t, env, "--config", flagCfg, "cluster", "list")
 	require.NoError(t, err)
-	assert.Equal(t, "account-from-flag", capturedAccountID)
+
+	req, ok := env.Server.ListClustersCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "account-from-flag", req.GetAccountId())
 }
 
 // TestConfigLoad_WithAccountIDTakesPrecedence verifies that WithAccountID (Set)
 // takes precedence over a config file loaded via --config (Set > config file).
 func TestConfigLoad_WithAccountIDTakesPrecedence(t *testing.T) {
 	env := testutil.NewTestEnv(t, testutil.WithAccountID("explicit-id"))
-	t.Cleanup(env.Cleanup)
 
 	cfgPath := testutil.WriteConfigFile(t, t.TempDir(), map[string]any{
 		"account_id": "account-from-file",
 	})
 
-	var capturedAccountID string
-	env.Server.ListClustersFunc = func(_ context.Context, req *clusterv1.ListClustersRequest) (*clusterv1.ListClustersResponse, error) {
-		capturedAccountID = req.GetAccountId()
-		return &clusterv1.ListClustersResponse{}, nil
-	}
+	env.Server.ListClustersCalls.Returns(&clusterv1.ListClustersResponse{}, nil)
 
 	_, _, err := testutil.Exec(t, env, "--config", cfgPath, "cluster", "list")
 	require.NoError(t, err)
-	assert.Equal(t, "explicit-id", capturedAccountID)
+
+	req, ok := env.Server.ListClustersCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "explicit-id", req.GetAccountId())
 }

internal/cmd/backup/create_test.go

@@ -1,7 +1,6 @@
 package backup_test
 
 import (
-	"context"
 	"encoding/json"
 	"testing"
 
@@ -15,50 +14,47 @@ import (
 
 func TestBackupCreate_Success(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
-
-	env.BackupServer.CreateBackupFunc = func(_ context.Context, req *backupv1.CreateBackupRequest) (*backupv1.CreateBackupResponse, error) {
-		assert.Equal(t, "test-account-id", req.GetBackup().GetAccountId())
-		assert.Equal(t, "cluster-abc", req.GetBackup().GetClusterId())
-		return &backupv1.CreateBackupResponse{
-			Backup: &backupv1.Backup{Id: "backup-new", ClusterId: "cluster-abc"},
-		}, nil
-	}
+
+	env.BackupServer.CreateBackupCalls.Returns(&backupv1.CreateBackupResponse{
+		Backup: &backupv1.Backup{Id: "backup-new", ClusterId: "cluster-abc"},
+	}, nil)
 
 	stdout, _, err := testutil.Exec(t, env, "backup", "create", "--cluster-id=cluster-abc", "--retention-days=7")
 	require.NoError(t, err)
 	assert.Contains(t, stdout, "backup-new")
 	assert.Contains(t, stdout, "cluster-abc")
+
+	req, ok := env.BackupServer.CreateBackupCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "test-account-id", req.GetBackup().GetAccountId())
+	assert.Equal(t, "cluster-abc", req.GetBackup().GetClusterId())
 }
 
 func TestBackupCreate_WithRetention(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	var capturedRetention int64
-	env.BackupServer.CreateBackupFunc = func(_ context.Context, req *backupv1.CreateBackupRequest) (*backupv1.CreateBackupResponse, error) {
-		if req.GetBackup().GetRetentionPeriod() != nil {
-			capturedRetention = int64(req.GetBackup().GetRetentionPeriod().AsDuration().Hours()) / 24
-		}
-		return &backupv1.CreateBackupResponse{
-			Backup: &backupv1.Backup{Id: "backup-ret", ClusterId: "cluster-abc"},
-		}, nil
-	}
+	env.BackupServer.CreateBackupCalls.Returns(&backupv1.CreateBackupResponse{
+		Backup: &backupv1.Backup{Id: "backup-ret", ClusterId: "cluster-abc"},
+	}, nil)
 
 	_, _, err := testutil.Exec(t, env, "backup", "create", "--cluster-id=cluster-abc", "--retention-days=7")
 	require.NoError(t, err)
+
+	req, ok := env.BackupServer.CreateBackupCalls.Last()
+	require.True(t, ok)
+	var capturedRetention int64
+	if req.GetBackup().GetRetentionPeriod() != nil {
+		capturedRetention = int64(req.GetBackup().GetRetentionPeriod().AsDuration().Hours()) / 24
+	}
 	assert.Equal(t, int64(7), capturedRetention)
 }
 
 func TestBackupCreate_JSONOutput(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.CreateBackupFunc = func(_ context.Context, _ *backupv1.CreateBackupRequest) (*backupv1.CreateBackupResponse, error) {
-		return &backupv1.CreateBackupResponse{
-			Backup: &backupv1.Backup{Id: "backup-json", ClusterId: "cluster-123"},
-		}, nil
-	}
+	env.BackupServer.CreateBackupCalls.Returns(&backupv1.CreateBackupResponse{
+		Backup: &backupv1.Backup{Id: "backup-json", ClusterId: "cluster-123"},
+	}, nil)
 
 	stdout, _, err := testutil.Exec(t, env, "backup", "create", "--cluster-id=cluster-123", "--retention-days=7", "--json")
 	require.NoError(t, err)
@@ -73,27 +69,22 @@ func TestBackupCreate_JSONOutput(t *testing.T) {
 
 func TestBackupCreate_InvalidRetention(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
 	_, _, err := testutil.Exec(t, env, "backup", "create", "--cluster-id=cluster-abc", "--retention-days=0")
 	require.Error(t, err)
 }
 
 func TestBackupCreate_MissingClusterID(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
 	_, _, err := testutil.Exec(t, env, "backup", "create")
 	require.Error(t, err)
 }
 
 func TestBackupCreate_APIError(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.CreateBackupFunc = func(_ context.Context, _ *backupv1.CreateBackupRequest) (*backupv1.CreateBackupResponse, error) {
-		return nil, assert.AnError
-	}
+	env.BackupServer.CreateBackupCalls.Returns(nil, assert.AnError)
 
 	_, _, err := testutil.Exec(t, env, "backup", "create", "--cluster-id=cluster-abc", "--retention-days=7")
 	require.Error(t, err)

internal/cmd/backup/delete_test.go

@@ -1,7 +1,6 @@
 package backup_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -14,29 +13,24 @@ import (
 
 func TestBackupDelete_WithForce(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	var capturedBackupID string
-	env.BackupServer.DeleteBackupFunc = func(_ context.Context, req *backupv1.DeleteBackupRequest) (*backupv1.DeleteBackupResponse, error) {
-		assert.Equal(t, "test-account-id", req.GetAccountId())
-		capturedBackupID = req.GetBackupId()
-		return &backupv1.DeleteBackupResponse{}, nil
-	}
+	env.BackupServer.DeleteBackupCalls.Returns(&backupv1.DeleteBackupResponse{}, nil)
 
 	stdout, _, err := testutil.Exec(t, env, "backup", "delete", "backup-abc", "--force")
 	require.NoError(t, err)
-	assert.Equal(t, "backup-abc", capturedBackupID)
+
+	req, ok := env.BackupServer.DeleteBackupCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "test-account-id", req.GetAccountId())
+	assert.Equal(t, "backup-abc", req.GetBackupId())
 	assert.Contains(t, stdout, "backup-abc")
 	assert.Contains(t, stdout, "deleted")
 }
 
 func TestBackupDelete_APIError(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.DeleteBackupFunc = func(_ context.Context, _ *backupv1.DeleteBackupRequest) (*backupv1.DeleteBackupResponse, error) {
-		return nil, assert.AnError
-	}
+	env.BackupServer.DeleteBackupCalls.Returns(nil, assert.AnError)
 
 	_, _, err := testutil.Exec(t, env, "backup", "delete", "backup-abc", "--force")
 	require.Error(t, err)

internal/cmd/backup/describe_test.go

@@ -1,7 +1,6 @@
 package backup_test
 
 import (
-	"context"
 	"encoding/json"
 	"testing"
 
@@ -16,39 +15,36 @@ import (
 
 func TestBackupDescribe_TextOutput(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.GetBackupFunc = func(_ context.Context, req *backupv1.GetBackupRequest) (*backupv1.GetBackupResponse, error) {
-		assert.Equal(t, "test-account-id", req.GetAccountId())
-		assert.Equal(t, "backup-abc", req.GetBackupId())
-		return &backupv1.GetBackupResponse{
-			Backup: &backupv1.Backup{
-				Id:        "backup-abc",
-				Name:      "my-backup",
-				ClusterId: "cluster-123",
-				Status:    backupv1.BackupStatus_BACKUP_STATUS_SUCCEEDED,
-				CreatedAt: timestamppb.Now(),
-			},
-		}, nil
-	}
+	env.BackupServer.GetBackupCalls.Returns(&backupv1.GetBackupResponse{
+		Backup: &backupv1.Backup{
+			Id:        "backup-abc",
+			Name:      "my-backup",
+			ClusterId: "cluster-123",
+			Status:    backupv1.BackupStatus_BACKUP_STATUS_SUCCEEDED,
+			CreatedAt: timestamppb.Now(),
+		},
+	}, nil)
 
 	stdout, _, err := testutil.Exec(t, env, "backup", "describe", "backup-abc")
 	require.NoError(t, err)
 	assert.Contains(t, stdout, "backup-abc")
 	assert.Contains(t, stdout, "my-backup")
 	assert.Contains(t, stdout, "cluster-123")
 	assert.Contains(t, stdout, "SUCCEEDED")
+
+	req, ok := env.BackupServer.GetBackupCalls.Last()
+	require.True(t, ok)
+	assert.Equal(t, "test-account-id", req.GetAccountId())
+	assert.Equal(t, "backup-abc", req.GetBackupId())
 }
 
 func TestBackupDescribe_JSONOutput(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.GetBackupFunc = func(_ context.Context, _ *backupv1.GetBackupRequest) (*backupv1.GetBackupResponse, error) {
-		return &backupv1.GetBackupResponse{
-			Backup: &backupv1.Backup{Id: "backup-json", ClusterId: "cluster-xyz"},
-		}, nil
-	}
+	env.BackupServer.GetBackupCalls.Returns(&backupv1.GetBackupResponse{
+		Backup: &backupv1.Backup{Id: "backup-json", ClusterId: "cluster-xyz"},
+	}, nil)
 
 	stdout, _, err := testutil.Exec(t, env, "backup", "describe", "backup-json", "--json")
 	require.NoError(t, err)
@@ -64,11 +60,8 @@ func TestBackupDescribe_JSONOutput(t *testing.T) {
 
 func TestBackupDescribe_APIError(t *testing.T) {
 	env := testutil.NewTestEnv(t)
-	t.Cleanup(env.Cleanup)
 
-	env.BackupServer.GetBackupFunc = func(_ context.Context, _ *backupv1.GetBackupRequest) (*backupv1.GetBackupResponse, error) {
-		return nil, assert.AnError
-	}
+	env.BackupServer.GetBackupCalls.Returns(nil, assert.AnError)
 
 	_, _, err := testutil.Exec(t, env, "backup", "describe", "backup-abc")
 	require.Error(t, err)