shudao-main/shudao-chat-go/tests/add_data.go

package tests

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// OSS配置信息
var (
	ossBucket    = "gdsc-ai-aqzs"
	ossAccessKey = "fnyfi2f368pbic74d8ll"
	ossSecretKey = "jgqwk7sirqlz2602x2k7yx2eor0vii19wah6ywlv"
	ossEndpoint  = "http://172.16.17.52:8060"
	ossRegion    = "us-east-1"
)

// 简单测试函数
func TestSimple(t *testing.T) {
	t.Logf("这是一个简单测试")
	t.Logf("测试日志输出")
}

// 批量上传图片测试函数
func TestBatchUploadImages(t *testing.T) {
	t.Logf("=== 开始批量上传图片测试 ===")

	// 先测试一个简单的日志输出
	t.Logf("测试日志输出正常")

	// 设置图片文件夹路径
	imageFolder := "C:/Users/allen/Desktop/隐患识别图/正确图片"
	t.Logf("目标文件夹: %s", imageFolder)

	// 检查文件夹是否存在
	if _, err := os.Stat(imageFolder); os.IsNotExist(err) {
		t.Fatalf("文件夹不存在: %s", imageFolder)
	}

	t.Logf("文件夹存在，开始扫描...")

	// 获取所有图片文件
	imageFiles, err := getImageFiles(imageFolder)
	if err != nil {
		t.Fatalf("获取图片文件失败: %v", err)
	}

	t.Logf("找到 %d 个图片文件", len(imageFiles))

	// 创建S3会话
	sess, err := createS3Session()
	if err != nil {
		t.Fatalf("创建S3会话失败: %v", err)
	}
	s3Client := s3.New(sess)

	// 处理每个图片文件（只处理前5个文件进行测试）
	successCount := 0
	errorCount := 0

	// 限制处理文件数量
	maxFiles := 5
	if len(imageFiles) > maxFiles {
		imageFiles = imageFiles[:maxFiles]
		t.Logf("为了测试，只处理前 %d 个文件", maxFiles)
	}

	for i, imageFile := range imageFiles {
		t.Logf("处理文件 %d/%d: %s", i+1, len(imageFiles), filepath.Base(imageFile))

		// 提取文件名中的序号
		fileID, err := extractIDFromFilename(filepath.Base(imageFile))
		if err != nil {
			t.Logf("  错误: 无法提取序号 - %v", err)
			errorCount++
			continue
		}

		t.Logf("  提取的序号: %d", fileID)

		// 暂时跳过数据库检查，直接测试OSS上传
		t.Logf("  跳过数据库检查，直接测试OSS上传")

		// 上传图片到OSS
		imageURL, err := uploadImageToOSS(s3Client, imageFile, fileID)
		if err != nil {
			t.Logf("  错误: 上传图片失败 - %v", err)
			errorCount++
			continue
		}

		t.Logf("  图片上传成功: %s", imageURL)

		// 暂时跳过数据库更新
		t.Logf("  跳过数据库更新")
		successCount++

		// 添加延迟避免请求过于频繁
		time.Sleep(100 * time.Millisecond)
	}

	t.Logf("处理完成！")
	t.Logf("成功处理: %d 个文件", successCount)
	t.Logf("处理失败: %d 个文件", errorCount)
}

// 获取文件夹中的所有图片文件
func getImageFiles(folderPath string) ([]string, error) {
	var imageFiles []string

	// 支持的图片格式
	imageExts := map[string]bool{
		".jpg":  true,
		".jpeg": true,
		".png":  true,
		".gif":  true,
		".bmp":  true,
		".webp": true,
		".tiff": true,
		".svg":  true,
		".ico":  true,
	}

	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			ext := strings.ToLower(filepath.Ext(path))
			if imageExts[ext] {
				imageFiles = append(imageFiles, path)
			}
		}

		return nil
	})

	return imageFiles, err
}

// 从文件名中提取序号
func extractIDFromFilename(filename string) (int, error) {
	// 移除文件扩展名
	nameWithoutExt := strings.TrimSuffix(filename, filepath.Ext(filename))

	// 使用正则表达式提取开头的数字
	re := regexp.MustCompile(`^(\d+)`)
	matches := re.FindStringSubmatch(nameWithoutExt)

	if len(matches) < 2 {
		return 0, fmt.Errorf("文件名中未找到数字序号: %s", filename)
	}

	id, err := strconv.Atoi(matches[1])
	if err != nil {
		return 0, fmt.Errorf("无法解析序号: %s", matches[1])
	}

	return id, nil
}

// 创建S3会话
func createS3Session() (*session.Session, error) {
	s3Config := &aws.Config{
		Credentials:      credentials.NewStaticCredentials(ossAccessKey, ossSecretKey, ""),
		Endpoint:         aws.String(ossEndpoint),
		Region:           aws.String(ossRegion),
		S3ForcePathStyle: aws.Bool(true),
	}

	sess, err := session.NewSession(s3Config)
	if err != nil {
		return nil, err
	}

	// 验证凭据
	_, err = sess.Config.Credentials.Get()
	if err != nil {
		return nil, fmt.Errorf("凭据验证失败: %v", err)
	}

	return sess, nil
}

// 上传图片到OSS
func uploadImageToOSS(s3Client *s3.S3, imagePath string, fileID int) (string, error) {
	// 打开图片文件
	file, err := os.Open(imagePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 生成文件名
	ext := filepath.Ext(imagePath)
	fileName := fmt.Sprintf("batch_upload/%d_%d%s",
		time.Now().Unix(), fileID, ext)

	// 读取文件内容
	fileBytes, err := io.ReadAll(file)
	if err != nil {
		return "", err
	}

	// 上传到S3
	_, err = s3Client.PutObject(&s3.PutObjectInput{
		Bucket: aws.String(ossBucket),
		Key:    aws.String(fileName),
		Body:   aws.ReadSeekCloser(strings.NewReader(string(fileBytes))),
		ACL:    aws.String("public-read"),
	})

	if err != nil {
		return "", err
	}

	// 生成访问URL
	imageURL := fmt.Sprintf("%s/%s/%s", ossEndpoint, ossBucket, fileName)
	return imageURL, nil
}