gotidb/pkg/storage/engine_test.go

package storage

import (
	"context"
	"os"
	"path/filepath"
	"testing"
	"time"

	"git.pyer.club/kingecg/gotidb/pkg/model"
)

func TestMemoryEngine(t *testing.T) {
	// 创建内存存储引擎
	engine := NewMemoryEngine()

	// 创建测试数据
	id := model.DataPointID{
		DeviceID:   "test-device",
		MetricCode: "temperature",
		Labels: map[string]string{
			"location": "room1",
		},
	}

	now := time.Now()
	value := model.DataValue{
		Timestamp: now,
		Value:     25.5,
	}

	// 测试写入
	t.Run("Write", func(t *testing.T) {
		err := engine.Write(context.Background(), id, value)
		if err != nil {
			t.Errorf("Write() error = %v", err)
		}
	})

	// 测试读取最新值
	t.Run("ReadLatest", func(t *testing.T) {
		latest, err := engine.ReadLatest(context.Background(), id)
		if err != nil {
			t.Errorf("ReadLatest() error = %v", err)
		}

		if !latest.Timestamp.Equal(value.Timestamp) {
			t.Errorf("ReadLatest() timestamp = %v, want %v", latest.Timestamp, value.Timestamp)
		}

		if latest.Value != value.Value {
			t.Errorf("ReadLatest() value = %v, want %v", latest.Value, value.Value)
		}
	})

	// 测试批量写入
	t.Run("BatchWrite", func(t *testing.T) {
		id2 := model.DataPointID{
			DeviceID:   "test-device",
			MetricCode: "humidity",
			Labels: map[string]string{
				"location": "room1",
			},
		}

		value2 := model.DataValue{
			Timestamp: now,
			Value:     60.0,
		}

		batch := []struct {
			ID    model.DataPointID
			Value model.DataValue
		}{
			{
				ID:    id,
				Value: value,
			},
			{
				ID:    id2,
				Value: value2,
			},
		}

		err := engine.BatchWrite(context.Background(), batch)
		if err != nil {
			t.Errorf("BatchWrite() error = %v", err)
		}

		// 验证批量写入的数据
		latest, err := engine.ReadLatest(context.Background(), id2)
		if err != nil {
			t.Errorf("ReadLatest() after BatchWrite error = %v", err)
		}

		if !latest.Timestamp.Equal(value2.Timestamp) {
			t.Errorf("ReadLatest() after BatchWrite timestamp = %v, want %v", latest.Timestamp, value2.Timestamp)
		}

		if latest.Value != value2.Value {
			t.Errorf("ReadLatest() after BatchWrite value = %v, want %v", latest.Value, value2.Value)
		}
	})

	// 测试读取所有值
	t.Run("ReadAll", func(t *testing.T) {
		// 写入多个值
		for i := 1; i <= 5; i++ {
			newValue := model.DataValue{
				Timestamp: now.Add(time.Duration(i) * time.Minute),
				Value:     25.5 + float64(i),
			}
			err := engine.Write(context.Background(), id, newValue)
			if err != nil {
				t.Errorf("Write() for ReadAll error = %v", err)
			}
		}

		// 读取所有值
		values, err := engine.ReadAll(context.Background(), id)
		if err != nil {
			t.Errorf("ReadAll() error = %v", err)
		}

		// 验证读取的值数量
		if len(values) != 6 { // 初始值 + 5个新值
			t.Errorf("ReadAll() returned %v values, want %v", len(values), 6)
		}

		// 验证值是按时间顺序排列的
		for i := 1; i < len(values); i++ {
			if values[i].Timestamp.Before(values[i-1].Timestamp) {
				t.Errorf("ReadAll() values not in chronological order")
			}
		}
	})

	// 测试读取持续时间内的值
	t.Run("ReadDuration", func(t *testing.T) {
		// 设置时间范围
		from := now.Add(1 * time.Minute)
		to := now.Add(3 * time.Minute)

		// 读取指定时间范围内的值
		values, err := engine.ReadDuration(context.Background(), id, from, to)
		if err != nil {
			t.Errorf("ReadDuration() error = %v", err)
		}

		// 验证读取的值数量
		if len(values) != 3 { // 1分钟、2分钟和3分钟的值
			t.Errorf("ReadDuration() returned %v values, want %v", len(values), 3)
		}

		// 验证所有值都在指定的时间范围内
		for _, v := range values {
			if v.Timestamp.Before(from) || v.Timestamp.After(to) {
				t.Errorf("ReadDuration() returned value with timestamp %v outside range [%v, %v]", v.Timestamp, from, to)
			}
		}
	})
}

func TestPersistence(t *testing.T) {
	// 创建临时目录
	tempDir, err := os.MkdirTemp("", "gotidb-test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tempDir)

	// 创建内存存储引擎
	engine := NewMemoryEngine()

	// 启用WAL持久化
	persistenceConfig := PersistenceConfig{
		Type:      PersistenceTypeWAL,
		Directory: tempDir,
		SyncEvery: 1, // 每次写入都同步
	}

	err = engine.EnablePersistence(persistenceConfig)
	if err != nil {
		t.Fatalf("EnablePersistence() error = %v", err)
	}

	// 创建测试数据
	id := model.DataPointID{
		DeviceID:   "test-device",
		MetricCode: "temperature",
		Labels: map[string]string{
			"location": "room1",
		},
	}

	now := time.Now()
	value := model.DataValue{
		Timestamp: now,
		Value:     25.5,
	}

	// 写入数据
	err = engine.Write(context.Background(), id, value)
	if err != nil {
		t.Errorf("Write() with persistence error = %v", err)
	}

	// 关闭引擎
	err = engine.Close()
	if err != nil {
		t.Errorf("Close() error = %v", err)
	}

	// 检查WAL文件是否存在
	walFiles, err := filepath.Glob(filepath.Join(tempDir, "*.wal"))
	if err != nil {
		t.Errorf("Failed to list WAL files: %v", err)
	}
	if len(walFiles) == 0 {
		t.Errorf("No WAL files found after persistence")
	}

	// 创建新的引擎并从WAL恢复
	newEngine := NewMemoryEngine()
	err = newEngine.EnablePersistence(persistenceConfig)
	if err != nil {
		t.Fatalf("EnablePersistence() for new engine error = %v", err)
	}

	// 读取恢复后的数据
	latest, err := newEngine.ReadLatest(context.Background(), id)
	if err != nil {
		t.Errorf("ReadLatest() after recovery error = %v", err)
	}

	// 验证恢复的数据
	if latest.Value != value.Value {
		t.Errorf("ReadLatest() after recovery value = %v, want %v", latest.Value, value.Value)
	}

	// 关闭新引擎
	err = newEngine.Close()
	if err != nil {
		t.Errorf("Close() new engine error = %v", err)
	}
}