gobridge/session_test.go

package gobridge

import (
	"context"
	"fmt"
	"reflect"
	"testing"
)

// sessionResult 对应 example/worker.py 里 session_result 返回的
// {"value": int, "steps": [int, ...]}。
type sessionResult struct {
	Value int   `json:"value"`
	Steps []int `json:"steps"`
}

// parseWorkerID 从 "[worker N] ..." 格式的字符串里解析出 worker 编号，
// example/worker.py 里 session_init/global_increment/global_get 都是这个格式。
func parseWorkerID(t *testing.T, msg string) int {
	t.Helper()
	var id int
	if _, err := fmt.Sscanf(msg, "[worker %d]", &id); err != nil {
		t.Fatalf("parse worker id from %q: %v", msg, err)
	}
	return id
}

// TestNewSessionIsolation 覆盖 example/main.go 的 demoSession 演示过的
// NewSession 功能：不同 session 之间状态互相隔离；同一个 worker 进程内的
// session 共享进程级全局变量，不同 worker 之间完全独立。
func TestNewSessionIsolation(t *testing.T) {
	pool := newTestPool(t, WithWorkers(2))
	ctx := context.Background()

	// NewSession 内部按 pool 共享的原子计数器轮询：第一次 idx=1%2=1，
	// 第二次 idx=2%2=0，第三次 idx=3%2=1——sessA 和 sessC 落在同一个 worker，
	// sessB 落在另一个（用全新 pool，不掺杂其它调用，保证这个顺序确定）。
	sessA := NewSession(pool)
	sessB := NewSession(pool)
	sessC := NewSession(pool)

	if _, err := Invoke[string](ctx, sessA, "session_init", "A", 100); err != nil {
		t.Fatalf("sessA init: %v", err)
	}
	if _, err := Invoke[string](ctx, sessB, "session_init", "B", 200); err != nil {
		t.Fatalf("sessB init: %v", err)
	}
	if _, err := Invoke[string](ctx, sessC, "session_init", "C", 300); err != nil {
		t.Fatalf("sessC init: %v", err)
	}

	// 各自 step，验证会话状态互不干扰
	if v, err := Invoke[int](ctx, sessA, "session_step", "A", 10); err != nil || v != 110 {
		t.Fatalf("sessA step(+10): v=%d err=%v, want 110", v, err)
	}
	if v, err := Invoke[int](ctx, sessA, "session_step", "A", 5); err != nil || v != 115 {
		t.Fatalf("sessA step(+5): v=%d err=%v, want 115", v, err)
	}
	if v, err := Invoke[int](ctx, sessB, "session_step", "B", 50); err != nil || v != 250 {
		t.Fatalf("sessB step(+50): v=%d err=%v, want 250", v, err)
	}
	if v, err := Invoke[int](ctx, sessC, "session_step", "C", 99); err != nil || v != 399 {
		t.Fatalf("sessC step(+99): v=%d err=%v, want 399（与 sessA 同 worker 但状态独立）", v, err)
	}

	rA, err := Invoke[sessionResult](ctx, sessA, "session_result", "A")
	if err != nil {
		t.Fatalf("sessA result: %v", err)
	}
	if want := (sessionResult{Value: 115, Steps: []int{10, 5}}); !reflect.DeepEqual(rA, want) {
		t.Fatalf("sessA result: want %+v, got %+v", want, rA)
	}

	rB, err := Invoke[sessionResult](ctx, sessB, "session_result", "B")
	if err != nil {
		t.Fatalf("sessB result: %v", err)
	}
	if want := (sessionResult{Value: 250, Steps: []int{50}}); !reflect.DeepEqual(rB, want) {
		t.Fatalf("sessB result: want %+v, got %+v", want, rB)
	}

	rC, err := Invoke[sessionResult](ctx, sessC, "session_result", "C")
	if err != nil {
		t.Fatalf("sessC result: %v", err)
	}
	if want := (sessionResult{Value: 399, Steps: []int{99}}); !reflect.DeepEqual(rC, want) {
		t.Fatalf("sessC result: want %+v, got %+v", want, rC)
	}

	// 进程级全局变量：sessA 和 sessC 同一个 worker 进程，应该共享 _global_counter；
	// sessB 是独立进程，从 0 开始，不受影响。
	rawA, err := Invoke[string](ctx, sessA, "global_increment", 10)
	if err != nil {
		t.Fatalf("sessA global_increment: %v", err)
	}
	workerA := parseWorkerID(t, rawA)

	rawC, err := Invoke[string](ctx, sessC, "global_increment", 5)
	if err != nil {
		t.Fatalf("sessC global_increment: %v", err)
	}
	workerC := parseWorkerID(t, rawC)

	rawB, err := Invoke[string](ctx, sessB, "global_increment", 99)
	if err != nil {
		t.Fatalf("sessB global_increment: %v", err)
	}
	workerB := parseWorkerID(t, rawB)

	if workerA != workerC {
		t.Fatalf("want sessA/sessC 落在同一个 worker，got worker %d / %d", workerA, workerC)
	}
	if workerB == workerA {
		t.Fatalf("want sessB 落在独立的 worker，got 跟 sessA/sessC 一样都是 worker %d", workerB)
	}

	wantA := fmt.Sprintf("[worker %d] counter = 15", workerA) // sessA(+10) + sessC(+5) 共享
	if rawA2, err := Invoke[string](ctx, sessA, "global_get"); err != nil || rawA2 != wantA {
		t.Fatalf("sessA global_get: want %q, got %q (err=%v)", wantA, rawA2, err)
	}

	wantB := fmt.Sprintf("[worker %d] counter = 99", workerB) // 独立进程，只有自己的 +99
	if rawB != wantB {
		t.Fatalf("sessB global_increment: want %q, got %q", wantB, rawB)
	}
}

// TestStickyCtx 覆盖 example/main.go 的 demoSession 演示过的 StickyCtx 功能：
// 相同的亲和 key 无论调用多少次，都应该稳定路由到同一个 worker 进程。
func TestStickyCtx(t *testing.T) {
	pool := newTestPool(t, WithWorkers(2))
	ctx := StickyCtx(context.Background(), "sticky-key")

	var firstWorker int
	for i := 0; i < 4; i++ {
		msg, err := Invoke[string](ctx, pool, "session_init", fmt.Sprintf("aff-%d", i), i)
		if err != nil {
			t.Fatalf("Invoke #%d: %v", i, err)
		}
		worker := parseWorkerID(t, msg)
		if i == 0 {
			firstWorker = worker
			continue
		}
		if worker != firstWorker {
			t.Fatalf("StickyCtx 路由不稳定：第 1 次落在 worker %d，第 %d 次落在 worker %d", firstWorker, i+1, worker)
		}
	}
}