seaweedfs/weed/admin/plugin/testing/mock_plugin.go

package testing

import (
	"fmt"
	"sync"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// MockPlugin simulates a plugin worker for testing
type MockPlugin struct {
	ID              string
	Name            string
	Version         string
	ProtocolVersion string
	Capabilities    []*plugin_pb.JobTypeCapability

	// Configuration
	DetectionEnabled bool
	ExecutionEnabled bool
	FailureMode      string // "" = success, "detection_error", "execution_error"
	DetectionDelay   time.Duration
	ExecutionDelay   time.Duration

	// Tracking
	mu               sync.RWMutex
	detectedJobs     map[string]*plugin_pb.DetectedJob
	executedJobs     map[string]*JobExecution
	registrationTime time.Time
	callCount        map[string]int
	errors           []string
}

// JobExecution tracks execution of a job
type JobExecution struct {
	JobID           string
	Config          []*plugin_pb.ConfigFieldValue
	Status          string
	ProgressPercent int32
	Messages        []*plugin_pb.JobExecutionMessage
	StartTime       time.Time
	EndTime         time.Time
	ErrorInfo       *plugin_pb.JobFailed
	mu              sync.RWMutex
}

// NewMockPlugin creates a new mock plugin
func NewMockPlugin(id, name, version string) *MockPlugin {
	return &MockPlugin{
		ID:               id,
		Name:             name,
		Version:          version,
		ProtocolVersion:  "v1",
		Capabilities:     make([]*plugin_pb.JobTypeCapability, 0),
		DetectionEnabled: true,
		ExecutionEnabled: true,
		DetectionDelay:   100 * time.Millisecond,
		ExecutionDelay:   100 * time.Millisecond,
		detectedJobs:     make(map[string]*plugin_pb.DetectedJob),
		executedJobs:     make(map[string]*JobExecution),
		registrationTime: time.Now(),
		callCount:        make(map[string]int),
		errors:           make([]string, 0),
	}
}

// AddCapability adds a job type capability to the plugin
func (mp *MockPlugin) AddCapability(jobType string, canDetect, canExecute bool) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	capability := &plugin_pb.JobTypeCapability{
		JobType:    jobType,
		CanDetect:  canDetect,
		CanExecute: canExecute,
		Version:    "v1",
	}
	mp.Capabilities = append(mp.Capabilities, capability)
}

// GetRegistrationMessage returns the plugin registration message
func (mp *MockPlugin) GetRegistrationMessage() *plugin_pb.PluginMessage {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	return &plugin_pb.PluginMessage{
		Content: &plugin_pb.PluginMessage_Register{
			Register: &plugin_pb.PluginRegister{
				PluginId:        mp.ID,
				Name:            mp.Name,
				Version:         mp.Version,
				ProtocolVersion: mp.ProtocolVersion,
				Capabilities:    mp.Capabilities,
			},
		},
	}
}

// GetHeartbeatMessage returns a heartbeat message
func (mp *MockPlugin) GetHeartbeatMessage(pendingJobs, runningJobs int32, cpuUsage, memoryUsage float32) *plugin_pb.PluginMessage {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	uptime := int64(time.Since(mp.registrationTime).Seconds())

	return &plugin_pb.PluginMessage{
		Content: &plugin_pb.PluginMessage_Heartbeat{
			Heartbeat: &plugin_pb.PluginHeartbeat{
				PluginId:        mp.ID,
				Timestamp:       timestamppb.Now(),
				UptimeSeconds:   uptime,
				PendingJobs:     pendingJobs,
				CpuUsagePercent: cpuUsage,
				MemoryUsageMb:   memoryUsage,
			},
		},
	}
}

// SimulateDetection simulates job detection
func (mp *MockPlugin) SimulateDetection(jobType string, detectedCount int) ([]*plugin_pb.DetectedJob, error) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.callCount["detection"]++

	if !mp.DetectionEnabled {
		err := fmt.Errorf("detection disabled for plugin %s", mp.ID)
		mp.errors = append(mp.errors, err.Error())
		return nil, err
	}

	if mp.FailureMode == "detection_error" {
		err := fmt.Errorf("simulated detection error")
		mp.errors = append(mp.errors, err.Error())
		return nil, err
	}

	time.Sleep(mp.DetectionDelay)

	var jobs []*plugin_pb.DetectedJob
	now := time.Now()

	for i := 0; i < detectedCount; i++ {
		jobKey := fmt.Sprintf("detected-job-%s-%d-%d", jobType, now.Unix(), i)
		job := &plugin_pb.DetectedJob{
			JobKey:      jobKey,
			JobType:     jobType,
			Description: fmt.Sprintf("Detected %s job %d", jobType, i),
			Priority:    int64(10 - i),
			Metadata:    make(map[string]string),
		}

		jobKey2 := fmt.Sprintf("%s-%d", jobType, i)
		mp.detectedJobs[jobKey2] = job
		jobs = append(jobs, job)
	}

	return jobs, nil
}

// SimulateExecution simulates job execution
func (mp *MockPlugin) SimulateExecution(jobID, jobType string, config []*plugin_pb.ConfigFieldValue) (*JobExecution, error) {
	mp.mu.Lock()

	mp.callCount["execution"]++

	if !mp.ExecutionEnabled {
		err := fmt.Errorf("execution disabled for plugin %s", mp.ID)
		mp.errors = append(mp.errors, err.Error())
		mp.mu.Unlock()
		return nil, err
	}

	if mp.FailureMode == "execution_error" {
		err := fmt.Errorf("simulated execution error")
		mp.errors = append(mp.errors, err.Error())
		mp.mu.Unlock()
		return nil, err
	}

	execution := &JobExecution{
		JobID:     jobID,
		Config:    config,
		Status:    "running",
		Messages:  make([]*plugin_pb.JobExecutionMessage, 0),
		StartTime: time.Now(),
	}

	mp.executedJobs[jobID] = execution
	mp.mu.Unlock()

	// Simulate progress updates
	for progress := 0; progress <= 100; progress += 25 {
		time.Sleep(mp.ExecutionDelay)

		mp.mu.Lock()
		if exec, ok := mp.executedJobs[jobID]; ok {
			exec.mu.Lock()
			exec.ProgressPercent = int32(progress)

			msg := &plugin_pb.JobExecutionMessage{
				JobId: jobID,
				Content: &plugin_pb.JobExecutionMessage_Progress{
					Progress: &plugin_pb.JobProgress{
						ProgressPercent: int32(progress),
						CurrentStep:     fmt.Sprintf("Step %d", progress/25),
						StatusMessage:   fmt.Sprintf("Executing step at %d%%", progress),
						UpdatedAt:       timestamppb.Now(),
					},
				},
			}
			exec.Messages = append(exec.Messages, msg)
			exec.mu.Unlock()
		}
		mp.mu.Unlock()
	}

	mp.mu.Lock()
	if exec, ok := mp.executedJobs[jobID]; ok {
		exec.mu.Lock()
		exec.Status = "completed"
		exec.ProgressPercent = 100
		exec.EndTime = time.Now()

		completed := &plugin_pb.JobExecutionMessage{
			JobId: jobID,
			Content: &plugin_pb.JobExecutionMessage_JobCompleted{
				JobCompleted: &plugin_pb.JobCompleted{
					CompletedAt: timestamppb.Now(),
					Summary:     fmt.Sprintf("Job %s completed successfully", jobID),
					Output: map[string]string{
						"result": "success",
						"job_id": jobID,
					},
				},
			},
		}
		exec.Messages = append(exec.Messages, completed)
		exec.mu.Unlock()
	}
	mp.mu.Unlock()

	return execution, nil
}

// SimulateExecutionFailure simulates a job execution failure
func (mp *MockPlugin) SimulateExecutionFailure(jobID string, errorCode, errorMessage string, retryable bool) (*JobExecution, error) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	execution := &JobExecution{
		JobID:     jobID,
		Status:    "failed",
		Messages:  make([]*plugin_pb.JobExecutionMessage, 0),
		StartTime: time.Now(),
		EndTime:   time.Now(),
	}

	failedMsg := &plugin_pb.JobFailed{
		ErrorCode:    errorCode,
		ErrorMessage: errorMessage,
		Retryable:    retryable,
		FailedAt:     timestamppb.Now(),
		RetryCount:   0,
	}

	msg := &plugin_pb.JobExecutionMessage{
		JobId: jobID,
		Content: &plugin_pb.JobExecutionMessage_JobFailed{
			JobFailed: failedMsg,
		},
	}

	execution.Messages = append(execution.Messages, msg)
	execution.ErrorInfo = failedMsg
	mp.executedJobs[jobID] = execution

	return execution, nil
}

// GetExecutionMessages returns execution messages for a job
func (mp *MockPlugin) GetExecutionMessages(jobID string) []*plugin_pb.JobExecutionMessage {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	if exec, ok := mp.executedJobs[jobID]; ok {
		exec.mu.RLock()
		defer exec.mu.RUnlock()

		result := make([]*plugin_pb.JobExecutionMessage, len(exec.Messages))
		copy(result, exec.Messages)
		return result
	}

	return nil
}

// GetCallCount returns the number of times a method was called
func (mp *MockPlugin) GetCallCount(method string) int {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	return mp.callCount[method]
}

// GetErrors returns all recorded errors
func (mp *MockPlugin) GetErrors() []string {
	mp.mu.RLock()
	defer mp.mu.RUnlock()

	result := make([]string, len(mp.errors))
	copy(result, mp.errors)
	return result
}

// SetFailureMode sets the failure simulation mode
func (mp *MockPlugin) SetFailureMode(mode string) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.FailureMode = mode
}

// SetDetectionDelay sets the detection simulation delay
func (mp *MockPlugin) SetDetectionDelay(delay time.Duration) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.DetectionDelay = delay
}

// SetExecutionDelay sets the execution simulation delay
func (mp *MockPlugin) SetExecutionDelay(delay time.Duration) {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.ExecutionDelay = delay
}

// Reset clears all state
func (mp *MockPlugin) Reset() {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	mp.detectedJobs = make(map[string]*plugin_pb.DetectedJob)
	mp.executedJobs = make(map[string]*JobExecution)
	mp.callCount = make(map[string]int)
	mp.errors = make([]string, 0)
	mp.FailureMode = ""
}