ragflow/internal/ingestion/testutil/mocks.go

//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package testutil

import (
	"context"
	"errors"
	"time"

	doctype "ragflow/internal/deepdoc/parser/type"
	"ragflow/internal/entity"
)

// ──────────────────────────────────────────────────────────
// Mock DAO Implementations
// ──────────────────────────────────────────────────────────

// MockDocDAO is a mock implementation of docGetter interface.
type MockDocDAO struct {
	Docs map[string]*entity.Document
	Err  error
}

func (m *MockDocDAO) GetByID(id string) (*entity.Document, error) {
	if m.Err != nil {
		return nil, m.Err
	}
	d, ok := m.Docs[id]
	if !ok {
		return nil, errors.New("record not found")
	}
	return d, nil
}

// MockF2dDAO is a mock implementation of f2dGetter interface.
type MockF2dDAO struct {
	Mappings map[string][]*entity.File2Document
	Err      error
}

func (m *MockF2dDAO) GetByDocumentID(docID string) ([]*entity.File2Document, error) {
	if m.Err != nil {
		return nil, m.Err
	}
	return m.Mappings[docID], nil
}

// MockFileDAO is a mock implementation of fileGetter interface.
type MockFileDAO struct {
	Files map[string]*entity.File
	Err   error
}

func (m *MockFileDAO) GetByID(id string) (*entity.File, error) {
	if m.Err != nil {
		return nil, m.Err
	}
	f, ok := m.Files[id]
	if !ok {
		return nil, errors.New("not found")
	}
	return f, nil
}

// ──────────────────────────────────────────────────────────
// Mock Storage Implementation
// ──────────────────────────────────────────────────────────

// MockStorage is a complete mock implementation of storage.Storage interface.
type MockStorage struct {
	Data         map[string][]byte
	Err          error
	PutCallCount int
	PutBucket    string
	PutFnm       string
	PutData      []byte
}

// NewMockStorage creates a new MockStorage with the given data.
func NewMockStorage(data map[string][]byte) *MockStorage {
	return &MockStorage{Data: data}
}

// NewMockStorageWithError creates a new MockStorage that returns the given error.
func NewMockStorageWithError(err error) *MockStorage {
	return &MockStorage{Err: err}
}

func (m *MockStorage) Health() bool { return true }

func (m *MockStorage) Get(bucket, objectName string, tenantID ...string) ([]byte, error) {
	if m.Err != nil {
		return nil, m.Err
	}
	key := bucket + "/" + objectName
	if len(tenantID) > 0 && tenantID[0] != "" {
		key = tenantID[0] + "/" + key
	}
	d, ok := m.Data[key]
	if !ok {
		d, ok = m.Data[objectName]
	}
	return d, nil
}

func (m *MockStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
	if m.Err != nil {
		return m.Err
	}
	m.PutCallCount++
	m.PutBucket = bucket
	m.PutFnm = fnm
	m.PutData = binary
	return nil
}

func (m *MockStorage) Remove(bucket, fnm string, tenantID ...string) error { return nil }

func (m *MockStorage) ObjExist(bucket, fnm string, tenantID ...string) bool { return true }

func (m *MockStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
	return "", nil
}

func (m *MockStorage) BucketExists(bucket string) bool { return true }

func (m *MockStorage) RemoveBucket(bucket string) error { return nil }

func (m *MockStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool { return true }

func (m *MockStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool { return true }

func (m *MockStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
	var objects []string
	for key := range m.Data {
		objects = append(objects, key)
	}
	return objects, nil
}

func (m *MockStorage) Close() error { return nil }

// ──────────────────────────────────────────────────────────
// Mock PDF Parser Implementation
// ──────────────────────────────────────────────────────────

// MockPDFParser is a mock implementation of pdfParser interface.
type MockPDFParser struct {
	Sections []map[string]any
	Err      error
}

// NewMockPDFParser creates a new MockPDFParser.
func NewMockPDFParser(sections []map[string]any, err error) *MockPDFParser {
	return &MockPDFParser{Sections: sections, Err: err}
}

func (m *MockPDFParser) ParseWithDeepDoc(ctx context.Context, filename string, data []byte, config doctype.ParserConfig) ([]map[string]any, error) {
	return m.Sections, m.Err
}