package gateway import ( "context" "fmt" "sync" "github.com/seaweedfs/seaweedfs/weed/filer_client" "github.com/seaweedfs/seaweedfs/weed/mq/kafka/integration" "github.com/seaweedfs/seaweedfs/weed/mq/kafka/protocol" filer_pb "github.com/seaweedfs/seaweedfs/weed/pb/filer_pb" schema_pb "github.com/seaweedfs/seaweedfs/weed/pb/schema_pb" ) // mockRecord implements the SMQRecord interface for testing type mockRecord struct { key []byte value []byte timestamp int64 offset int64 } func (r *mockRecord) GetKey() []byte { return r.key } func (r *mockRecord) GetValue() []byte { return r.value } func (r *mockRecord) GetTimestamp() int64 { return r.timestamp } func (r *mockRecord) GetOffset() int64 { return r.offset } // mockSeaweedMQHandler is a stateful mock for unit testing without real SeaweedMQ type mockSeaweedMQHandler struct { mu sync.RWMutex topics map[string]*integration.KafkaTopicInfo records map[string]map[int32][]integration.SMQRecord // topic -> partition -> records offsets map[string]map[int32]int64 // topic -> partition -> next offset } func newMockSeaweedMQHandler() *mockSeaweedMQHandler { return &mockSeaweedMQHandler{ topics: make(map[string]*integration.KafkaTopicInfo), records: make(map[string]map[int32][]integration.SMQRecord), offsets: make(map[string]map[int32]int64), } } func (m *mockSeaweedMQHandler) TopicExists(topic string) bool { m.mu.RLock() defer m.mu.RUnlock() _, exists := m.topics[topic] return exists } func (m *mockSeaweedMQHandler) ListTopics() []string { m.mu.RLock() defer m.mu.RUnlock() topics := make([]string, 0, len(m.topics)) for topic := range m.topics { topics = append(topics, topic) } return topics } func (m *mockSeaweedMQHandler) CreateTopic(topic string, partitions int32) error { m.mu.Lock() defer m.mu.Unlock() if _, exists := m.topics[topic]; exists { return fmt.Errorf("topic already exists") } m.topics[topic] = &integration.KafkaTopicInfo{ Name: topic, Partitions: partitions, } return nil } func (m *mockSeaweedMQHandler) CreateTopicWithSchemas(name string, partitions int32, keyRecordType *schema_pb.RecordType, valueRecordType *schema_pb.RecordType) error { m.mu.Lock() defer m.mu.Unlock() if _, exists := m.topics[name]; exists { return fmt.Errorf("topic already exists") } m.topics[name] = &integration.KafkaTopicInfo{ Name: name, Partitions: partitions, } return nil } func (m *mockSeaweedMQHandler) DeleteTopic(topic string) error { m.mu.Lock() defer m.mu.Unlock() delete(m.topics, topic) return nil } func (m *mockSeaweedMQHandler) GetTopicInfo(topic string) (*integration.KafkaTopicInfo, bool) { m.mu.RLock() defer m.mu.RUnlock() info, exists := m.topics[topic] return info, exists } func (m *mockSeaweedMQHandler) ProduceRecord(topicName string, partitionID int32, key, value []byte) (int64, error) { m.mu.Lock() defer m.mu.Unlock() // Check if topic exists if _, exists := m.topics[topicName]; !exists { return 0, fmt.Errorf("topic does not exist: %s", topicName) } // Initialize partition records if needed if _, exists := m.records[topicName]; !exists { m.records[topicName] = make(map[int32][]integration.SMQRecord) m.offsets[topicName] = make(map[int32]int64) } // Get next offset offset := m.offsets[topicName][partitionID] m.offsets[topicName][partitionID]++ // Store record record := &mockRecord{ key: key, value: value, offset: offset, } m.records[topicName][partitionID] = append(m.records[topicName][partitionID], record) return offset, nil } func (m *mockSeaweedMQHandler) ProduceRecordValue(topicName string, partitionID int32, key []byte, recordValueBytes []byte) (int64, error) { return m.ProduceRecord(topicName, partitionID, key, recordValueBytes) } func (m *mockSeaweedMQHandler) GetStoredRecords(ctx context.Context, topic string, partition int32, fromOffset int64, maxRecords int) ([]integration.SMQRecord, error) { m.mu.RLock() defer m.mu.RUnlock() // Check if topic exists if _, exists := m.topics[topic]; !exists { return nil, fmt.Errorf("topic does not exist: %s", topic) } // Get partition records partitionRecords, exists := m.records[topic][partition] if !exists || len(partitionRecords) == 0 { return []integration.SMQRecord{}, nil } // Find records starting from fromOffset result := make([]integration.SMQRecord, 0, maxRecords) for _, record := range partitionRecords { if record.GetOffset() >= fromOffset { result = append(result, record) if len(result) >= maxRecords { break } } } return result, nil } func (m *mockSeaweedMQHandler) GetEarliestOffset(topic string, partition int32) (int64, error) { m.mu.RLock() defer m.mu.RUnlock() // Check if topic exists if _, exists := m.topics[topic]; !exists { return 0, fmt.Errorf("topic does not exist: %s", topic) } // Get partition records partitionRecords, exists := m.records[topic][partition] if !exists || len(partitionRecords) == 0 { return 0, nil } return partitionRecords[0].GetOffset(), nil } func (m *mockSeaweedMQHandler) GetLatestOffset(topic string, partition int32) (int64, error) { m.mu.RLock() defer m.mu.RUnlock() // Check if topic exists if _, exists := m.topics[topic]; !exists { return 0, fmt.Errorf("topic does not exist: %s", topic) } // Return next offset (latest + 1) if offsets, exists := m.offsets[topic]; exists { return offsets[partition], nil } return 0, nil } func (m *mockSeaweedMQHandler) WithFilerClient(streamingMode bool, fn func(filer_pb.SeaweedFilerClient) error) error { return fmt.Errorf("mock handler: not implemented") } func (m *mockSeaweedMQHandler) CreatePerConnectionBrokerClient() (*integration.BrokerClient, error) { // Return a minimal broker client that won't actually connect return nil, fmt.Errorf("mock handler: per-connection broker client not available in unit test mode") } func (m *mockSeaweedMQHandler) GetFilerClientAccessor() *filer_client.FilerClientAccessor { return nil } func (m *mockSeaweedMQHandler) GetBrokerAddresses() []string { return []string{"localhost:9092"} // Return a dummy broker address for unit tests } func (m *mockSeaweedMQHandler) Close() error { return nil } func (m *mockSeaweedMQHandler) SetProtocolHandler(h integration.ProtocolHandler) {} // NewMinimalTestHandler creates a minimal handler for unit testing // that won't actually process Kafka protocol requests func NewMinimalTestHandler() *protocol.Handler { return protocol.NewTestHandlerWithMock(newMockSeaweedMQHandler()) }