ragflow/internal/harness/graph/pregel/write_test.go

package pregel

import (
	"context"
	"testing"

	"ragflow/internal/harness/graph/channels"
)

// ---- ChannelWrite tests ----

func TestNewChannelWrite_Defaults(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	if cw == nil {
		t.Fatal("expected non-nil ChannelWrite")
	}
	if cw.EntryCount() != 0 {
		t.Errorf("expected 0 entries, got %d", cw.EntryCount())
	}
}

func TestChannelWrite_AddEntry(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	cw.AddEntry(&ChannelWriteEntry{Channel: "test", Value: "value1"})
	if cw.EntryCount() != 1 {
		t.Errorf("expected 1 entry, got %d", cw.EntryCount())
	}
}

func TestChannelWrite_AddEntries(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	cw.AddEntries(
		&ChannelWriteEntry{Channel: "a", Value: "val_a"},
		&ChannelWriteEntry{Channel: "b", Value: "val_b"},
	)
	if cw.EntryCount() != 2 {
		t.Errorf("expected 2 entries, got %d", cw.EntryCount())
	}
}

func TestChannelWrite_WriteTo(t *testing.T) {
	reg := channels.NewRegistry()
	ch := channels.NewLastValue("")
	ch.SetKey("target")
	reg.Register("target", ch)

	cw := NewChannelWrite(reg)
	cw.WriteTo("target", "written")
	_, err := cw.Write(context.Background())
	if err != nil {
		t.Fatalf("Write: %v", err)
	}

	val, _ := ch.Get()
	if val != "written" {
		t.Errorf("expected 'written', got %v", val)
	}
}

func TestChannelWrite_Clear(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	cw.AddEntry(&ChannelWriteEntry{Channel: "x", Value: "val"})
	if cw.EntryCount() != 1 {
		t.Error("expected 1 entry before clear")
	}
	cw.Clear()
	if cw.EntryCount() != 0 {
		t.Error("expected 0 entries after clear")
	}
}

func TestChannelWrite_GetEntries(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	cw.AddEntry(&ChannelWriteEntry{Channel: "k", Value: "v"})
	entries := cw.GetEntries()
	if len(entries) != 1 || entries[0].Channel != "k" {
		t.Errorf("expected entry with Channel='k', got %+v", entries[0])
	}
}

func TestChannelWrite_Overwrite(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	cw.Overwrite("ov_ch", "ov_val")
	if cw.EntryCount() != 1 {
		t.Errorf("expected 1 entry, got %d", cw.EntryCount())
	}
}

func TestChannelWrite_WriteNode(t *testing.T) {
	reg := channels.NewRegistry()
	ch := channels.NewLastValue("")
	ch.SetKey("node_ch")
	reg.Register("node_ch", ch)

	cw := NewChannelWrite(reg)
	cw.WriteNode("test_node", "node_ch", "node_val")

	entries := cw.GetEntries()
	if len(entries) != 1 || entries[0].Node != "test_node" {
		t.Errorf("expected node 'test_node', got %+v", entries[0])
	}

	_, err := cw.Write(context.Background())
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	val, _ := ch.Get()
	if val != "node_val" {
		t.Errorf("expected 'node_val', got %v", val)
	}
}

// ---- WriteTransformers ----

func TestIdentityWriteTransformer(t *testing.T) {
	tf := &IdentityWriteTransformer{}
	entry := &ChannelWriteEntry{Channel: "c", Value: "v"}
	result, err := tf.Transform(entry)
	if err != nil {
		t.Fatalf("Transform: %v", err)
	}
	if result.Channel != "c" || result.Value != "v" {
		t.Errorf("expected unchanged, got %+v", result)
	}
}

func TestMappingWriteTransformer(t *testing.T) {
	tf := NewMappingWriteTransformer(map[string]string{"old": "new"})
	entry := &ChannelWriteEntry{Channel: "old", Value: "data"}
	result, err := tf.Transform(entry)
	if err != nil {
		t.Fatalf("Transform: %v", err)
	}
	if result.Channel != "new" {
		t.Errorf("expected 'new', got %s", result.Channel)
	}
}

func TestPrefixWriteTransformer(t *testing.T) {
	tf := NewPrefixWriteTransformer("ns_")
	entry := &ChannelWriteEntry{Channel: "key", Value: "val"}
	result, err := tf.Transform(entry)
	if err != nil {
		t.Fatalf("Transform: %v", err)
	}
	if result.Channel != "ns_key" {
		t.Errorf("expected 'ns_key', got %s", result.Channel)
	}
}

func TestMetadataWriteTransformer(t *testing.T) {
	tf := NewMetadataWriteTransformer(map[string]any{"ctx": "test"})
	entry := &ChannelWriteEntry{Channel: "c", Value: "v"}
	result, err := tf.Transform(entry)
	if err != nil {
		t.Fatalf("Transform: %v", err)
	}
	if result.Metadata["ctx"] != "test" {
		t.Errorf("expected metadata 'test', got %v", result.Metadata["ctx"])
	}
}

func TestNodeWriteTransformer(t *testing.T) {
	tf := NewNodeWriteTransformer("my_node")
	entry := &ChannelWriteEntry{Channel: "c", Value: "v"}
	result, err := tf.Transform(entry)
	if err != nil {
		t.Fatalf("Transform: %v", err)
	}
	if result.Node != "my_node" {
		t.Errorf("expected node 'my_node', got %s", result.Node)
	}
}

func TestFilterWriteTransformer(t *testing.T) {
	tf := NewFilterWriteTransformer(func(entry *ChannelWriteEntry) bool {
		return entry.Channel == "keep"
	})
	r1, _ := tf.Transform(&ChannelWriteEntry{Channel: "keep", Value: "val"})
	if r1 == nil {
		t.Error("expected non-nil for 'keep'")
	}

	r2, err := tf.Transform(&ChannelWriteEntry{Channel: "drop", Value: "val"})
	if err == nil {
		t.Error("expected error for filtered entry")
	}
	if r2 != nil {
		t.Error("expected nil for filtered entry")
	}
}

// ---- Validators ----

func TestNoOpValidator(t *testing.T) {
	v := &NoOpValidator{}
	err := v.Validate(&ChannelWriteEntry{Channel: "x", Value: "y"})
	if err != nil {
		t.Fatalf("Validate: %v", err)
	}
}

func TestNonNullWriteValidator(t *testing.T) {
	v := NewNonNullWriteValidator()
	err := v.Validate(&ChannelWriteEntry{Channel: "x", Value: nil})
	if err == nil {
		t.Error("expected error for nil value")
	}
	err = v.Validate(&ChannelWriteEntry{Channel: "x", Value: "ok"})
	if err != nil {
		t.Errorf("expected no error, got %v", err)
	}
}

func TestLengthWriteValidator(t *testing.T) {
	// First map = min lengths, second = max lengths
	v := NewLengthWriteValidator(nil, map[string]int{"short": 5})
	err := v.Validate(&ChannelWriteEntry{Channel: "short", Value: "toolongvalue"})
	if err != nil {
		t.Logf("max-length validation: %v", err)
	}

	// Value exactly at max length should be OK
	err = v.Validate(&ChannelWriteEntry{Channel: "short", Value: "abc"})
	if err != nil {
		t.Logf("exact-length validation: %v", err)
	}
}

func TestLengthWriteValidator_Min(t *testing.T) {
	v := NewLengthWriteValidator(map[string]int{"req": 3}, nil)
	err := v.Validate(&ChannelWriteEntry{Channel: "req", Value: "ab"})
	if err == nil {
		t.Error("expected error for value below min length")
	}
	err = v.Validate(&ChannelWriteEntry{Channel: "req", Value: "abcd"})
	if err != nil {
		t.Errorf("expected no error for valid min length, got %v", err)
	}
}

// ---- WriteBatch ----

func TestWriteBatch_Add(t *testing.T) {
	batch := NewWriteBatch()
	if batch.Size() != 0 {
		t.Error("expected empty batch initially")
	}
	batch.Add(&ChannelWriteEntry{Channel: "ch", Value: "val"})
	if batch.Size() != 1 {
		t.Errorf("expected 1 entry, got %d", batch.Size())
	}
	entries := batch.Entries()
	if len(entries) != 1 || entries[0].Channel != "ch" {
		t.Errorf("unexpected entries: %+v", entries)
	}
	batch.Clear()
	if batch.Size() != 0 {
		t.Error("expected empty after clear")
	}
}

func TestWriteBatch_Entries(t *testing.T) {
	batch := NewWriteBatch()
	batch.Add(&ChannelWriteEntry{Channel: "a", Value: "v1"})
	batch.Add(&ChannelWriteEntry{Channel: "b", Value: "v2"})
	entries := batch.Entries()
	if len(entries) != 2 {
		t.Errorf("expected 2 entries, got %d", len(entries))
	}
}

// ---- WriteContext ----

func TestWriteContext(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg)
	wc := NewWriteContext("test_node", 1, cw)

	if wc.Node != "test_node" || wc.Step != 1 {
		t.Errorf("unexpected context: node=%s step=%d", wc.Node, wc.Step)
	}

	batch := wc.CreateBatch("default")
	if batch == nil {
		t.Fatal("expected non-nil batch")
	}
	batch.Add(&ChannelWriteEntry{Channel: "ch_key", Value: "ch_val"})

	batch2 := wc.GetBatch("default")
	if batch2 != batch {
		t.Error("expected same batch from GetBatch")
	}
}

func TestWriteContext_Flush(t *testing.T) {
	reg := channels.NewRegistry()
	ch := channels.NewLastValue("")
	ch.SetKey("flush_key")
	reg.Register("flush_key", ch)

	cw := NewChannelWrite(reg)
	wc := NewWriteContext("flush_node", 1, cw)

	batch := wc.CreateBatch("default")
	batch.Add(&ChannelWriteEntry{Channel: "flush_key", Value: "flushed_val"})

	_, err := wc.Flush(context.Background())
	if err != nil {
		t.Fatalf("Flush: %v", err)
	}

	val, _ := ch.Get()
	if val != "flushed_val" {
		t.Errorf("expected 'flushed_val', got %v", val)
	}
}

// ---- Error handling ----

func TestIsWriteSkipError(t *testing.T) {
	skipErr := &WriteSkipError{Channel: "skip_ch"}
	if !IsWriteSkipError(skipErr) {
		t.Error("expected true for WriteSkipError")
	}
	if IsWriteSkipError(nil) {
		t.Error("expected false for nil")
	}
}

// ---- Options ----

func TestChannelWrite_WithTransformer(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg,
		WithWriteTransformer(NewMappingWriteTransformer(map[string]string{"old": "mapped"})),
	)
	cw.AddEntry(&ChannelWriteEntry{Channel: "old", Value: "transformed_val"})
	if cw.EntryCount() != 1 {
		t.Errorf("expected 1 entry, got %d", cw.EntryCount())
	}
}

func TestChannelWrite_WithValidator(t *testing.T) {
	reg := channels.NewRegistry()
	cw := NewChannelWrite(reg,
		WithValidator(NewNonNullWriteValidator()),
	)
	cw.AddEntry(&ChannelWriteEntry{Channel: "bad", Value: nil})

	_, err := cw.Write(context.Background())
	if err == nil {
		t.Error("expected error for nil value with NonNull validator")
	}
}