gRPC Mocks and Fakes — Building Test Doubles for Reliable Service Testing

package main

//go:generate mockgen -source=threat/v1/threat_grpc.pb.go \
//    -destination=threat/v1/mock_threat.go \
//    -package=threatv1

import (
    "testing"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/status"
    "google.golang.org/protobuf/proto"
    
    threatv1 "path/to/threat/v1"
)

func TestReportThreat_RetryOnUnavailable(t *testing.T) {
    // Create mock
    mockClient := threatv1.NewMockThreatServiceClient(gomock.NewController(t))
    
    // Set up expectations: first two calls fail, third succeeds
    firstCall := mockClient.EXPECT().
        ReportThreat(gomock.Any(), gomock.Any()).
        Return(nil, status.Error(codes.Unavailable, "server restarting"))
    
    secondCall := mockClient.EXPECT().
        ReportThreat(gomock.Any(), gomock.Any()).
        Return(nil, status.Error(codes.Unavailable, "server restarting"))
    
    thirdCall := mockClient.EXPECT().
        ReportThreat(gomock.Any(), gomock.Any()).
        Return(&threatv1.ThreatResponse{
            ThreatId: "threat-123",
            Status:   threatv1.ThreatStatus_QUARANTINED,
        }, nil)
    
    // Run client under test
    client := NewThreatClient(mockClient, RetryConfig{
        MaxRetries: 3,
        BaseDelay:  10 * time.Millisecond,
    })
    
    resp, err := client.ReportThreat(context.Background(),
        &threatv1.ThreatRequest{Name: "Malware-X"})
    
    if err != nil {
        t.Fatalf("Expected success, got %v", err)
    }
    if resp.ThreatId != "threat-123" {
        t.Errorf("Expected threat-123, got %s", resp.ThreatId)
    }
    
    // Verify all expectations met
    gomock.AssertExpectations(t)
}

import unittest
from unittest.mock import MagicMock, patch, call
import grpc

class TestThreatClient(unittest.TestCase):
    def setUp(self):
        # Create mock stub
        self.mock_stub = MagicMock()
        self.client = ThreatClient(self.mock_stub)
    
    def test_report_threat_success(self):
        # Configure mock response
        self.mock_stub.ReportThreat.return_value = pb.ThreatResponse(
            threat_id="threat-123",
            status=pb.ThreatStatus.QUARANTINED,
        )
        
        # Call client
        result = self.client.report_threat(
            device_id="dev-001",
            threat_name="Ransomware-X",
        )
        
        # Verify result
        self.assertEqual(result.threat_id, "threat-123")
        
        # Verify exact call arguments
        self.mock_stub.ReportThreat.assert_called_once()
        call_args = self.mock_stub.ReportThreat.call_args[0][0]
        self.assertEqual(call_args.device_id, "dev-001")
        self.assertEqual(call_args.threat_name, "Ransomware-X")
    
    def test_retry_on_unavailable(self):
        # Mock raises UNAVAILABLE twice, then succeeds
        self.mock_stub.ReportThreat.side_effect = [
            grpc.RpcError(
                grpc.StatusCode.UNAVAILABLE,
                "Unavailable"),
            grpc.RpcError(
                grpc.StatusCode.UNAVAILABLE,
                "Unavailable"),
            pb.ThreatResponse(threat_id="threat-456"),
        ]
        
        result = self.client.report_threat(
            device_id="dev-002",
            threat_name="Trojan-Y",
        )
        
        # Verify 3 calls were made
        self.assertEqual(
            self.mock_stub.ReportThreat.call_count, 3)
        self.assertEqual(result.threat_id, "threat-456")
    
    def test_non_retryable_error(self):
        # INVALID_ARGUMENT should not be retried
        self.mock_stub.ReportThreat.side_effect = grpc.RpcError(
            grpc.StatusCode.INVALID_ARGUMENT,
            "Invalid device ID",
        )
        
        with self.assertRaises(RuntimeError):
            self.client.report_threat(
                device_id="", threat_name="Malware")
        
        # Verify only 1 call (no retry)
        self.mock_stub.ReportThret.assert_called_once()

class FakeThreatServer {
  constructor() {
    this.threats = [];
    this.shouldFail = false;
    this.failCount = 0;
    this.callHistory = [];
  }
  
  start() {
    this.server = new grpc.Server();
    this.server.addService(ThreatService.service, {
      ReportThreat: (call, callback) => {
        this.callHistory.push({
          method: 'ReportThreat',
          request: call.request,
          timestamp: Date.now(),
        });
        
        if (this.shouldFail && this.failCount > 0) {
          this.failCount--;
          callback({ code: grpc.status.UNAVAILABLE, details: 'Simulated failure' });
          return;
        }
        
        const threat = {
          threatId: `threat-${this.threats.length + 1}`,
          name: call.request.name,
          status: 'QUARANTINED',
        };
        this.threats.push(threat);
        callback(null, threat);
      },
      
      ListThreats: (call, callback) => {
        callback(null, { threats: this.threats });
      },
    });
    
    return new Promise((resolve, reject) => {
      this.server.bindAsync('0.0.0.0:0',
        grpc.ServerCredentials.createInsecure(),
        (err, port) => {
          if (err) return reject(err);
          this.port = port;
          this.server.start();
          resolve(port);
        }
      );
    });
  }
  
  stop() {
    this.server.forceShutdown();
  }
}

// Test using fake server
describe('ThreatClient with Fake Server', () => {
  it('should retry and succeed after failures', async () => {
    const fake = new FakeThreatServer();
    const port = await fake.start();
    fake.shouldFail = true;
    fake.failCount = 2;
    
    const client = new ThreatClient(`localhost:${port}`);
    const result = await client.reportThreat('dev-1', 'Malware');
    
    assert.strictEqual(result.threatId, 'threat-1');
    assert.strictEqual(fake.callHistory.length, 3);
    fake.stop();
  });
});