rdp-proxy/agents/rap-node-agent/internal/fabricproto/session_test.go

package fabricproto

import (
	"errors"
	"testing"
)

func TestSessionDequeuesByTrafficPriority(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 10, TrafficClassBulk)
	mustOpenStream(t, session, 20, TrafficClassInteractive)
	mustOpenStream(t, session, 30, TrafficClassControl)

	mustEnqueue(t, session, 10, "bulk")
	mustEnqueue(t, session, 20, "interactive")
	mustEnqueue(t, session, 30, "control")

	frame, ok := session.DequeueNext()
	if !ok || frame.StreamID != 30 {
		t.Fatalf("first frame = %+v ok=%v, want control stream 30", frame, ok)
	}
	frame, ok = session.DequeueNext()
	if !ok || frame.StreamID != 20 {
		t.Fatalf("second frame = %+v ok=%v, want interactive stream 20", frame, ok)
	}
	frame, ok = session.DequeueNext()
	if !ok || frame.StreamID != 10 {
		t.Fatalf("third frame = %+v ok=%v, want bulk stream 10", frame, ok)
	}
}

func TestSessionBulkQueueFullDoesNotBlockControlOrInteractive(t *testing.T) {
	session := NewSession(SessionConfig{
		ClassQueueCapacity: map[TrafficClass]int{
			TrafficClassBulk:        1,
			TrafficClassControl:     4,
			TrafficClassInteractive: 4,
		},
	})
	mustOpenStream(t, session, 1, TrafficClassBulk)
	mustOpenStream(t, session, 2, TrafficClassControl)
	mustOpenStream(t, session, 3, TrafficClassInteractive)

	mustEnqueue(t, session, 1, "bulk-1")
	if _, err := session.EnqueueData(1, []byte("bulk-2")); !errors.Is(err, ErrStreamQueueFull) {
		t.Fatalf("bulk enqueue error = %v, want %v", err, ErrStreamQueueFull)
	}
	mustEnqueue(t, session, 2, "control")
	mustEnqueue(t, session, 3, "interactive")

	first, ok := session.DequeueNext()
	if !ok || first.StreamID != 2 {
		t.Fatalf("first frame = %+v ok=%v, want control", first, ok)
	}
	second, ok := session.DequeueNext()
	if !ok || second.StreamID != 3 {
		t.Fatalf("second frame = %+v ok=%v, want interactive", second, ok)
	}

	snapshot := session.Snapshot()
	if snapshot.QueueFull != 1 || snapshot.FramesDropped != 1 {
		t.Fatalf("snapshot queue full/dropped = %d/%d, want 1/1", snapshot.QueueFull, snapshot.FramesDropped)
	}
	if snapshot.Streams[2].Enqueued != 1 || snapshot.Streams[3].Enqueued != 1 {
		t.Fatalf("control/interactive metrics = %+v / %+v", snapshot.Streams[2], snapshot.Streams[3])
	}
}

func TestSessionCreditBlocksOnlyOneStream(t *testing.T) {
	session := NewSession(SessionConfig{InitialStreamCredit: 1})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	mustOpenStream(t, session, 2, TrafficClassReliable)

	mustEnqueue(t, session, 1, "one")
	if _, err := session.EnqueueData(1, []byte("two")); !errors.Is(err, ErrStreamCreditExhausted) {
		t.Fatalf("credit error = %v, want %v", err, ErrStreamCreditExhausted)
	}
	mustEnqueue(t, session, 2, "other")

	snapshot := session.Snapshot()
	if snapshot.CreditExhausted != 1 || snapshot.Streams[1].CreditBlocked != 1 {
		t.Fatalf("credit metrics = %+v stream=%+v", snapshot, snapshot.Streams[1])
	}
	if snapshot.Streams[2].Enqueued != 1 {
		t.Fatalf("second stream metrics = %+v, want enqueued", snapshot.Streams[2])
	}
}

func TestSessionAddCreditAllowsMoreFrames(t *testing.T) {
	session := NewSession(SessionConfig{InitialStreamCredit: 1})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	mustEnqueue(t, session, 1, "one")
	if _, err := session.EnqueueData(1, []byte("blocked")); !errors.Is(err, ErrStreamCreditExhausted) {
		t.Fatalf("credit error = %v, want %v", err, ErrStreamCreditExhausted)
	}
	if err := session.AddCredit(1, 2); err != nil {
		t.Fatalf("add credit: %v", err)
	}
	mustEnqueue(t, session, 1, "two")

	snapshot := session.Snapshot()
	if snapshot.Streams[1].Credit != 1 || snapshot.Streams[1].Enqueued != 2 {
		t.Fatalf("stream metrics = %+v, want credit 1 and enqueued 2", snapshot.Streams[1])
	}
}

func TestSessionResetDropsOnlySelectedStream(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 1, TrafficClassBulk)
	mustOpenStream(t, session, 2, TrafficClassControl)
	mustEnqueue(t, session, 1, "bulk")
	mustEnqueue(t, session, 2, "control")

	if err := session.ResetStream(1); err != nil {
		t.Fatalf("reset stream: %v", err)
	}
	frame, ok := session.DequeueNext()
	if !ok || frame.StreamID != 2 {
		t.Fatalf("dequeued frame = %+v ok=%v, want control stream 2", frame, ok)
	}
	if _, ok := session.DequeueNext(); ok {
		t.Fatal("unexpected frame after reset dropped bulk queue")
	}

	snapshot := session.Snapshot()
	if snapshot.StreamsReset != 1 || snapshot.Streams[1].State != StreamStateReset || snapshot.Streams[2].Dequeued != 1 {
		t.Fatalf("snapshot = %+v", snapshot)
	}
}

func TestSessionAckUpdatesMetrics(t *testing.T) {
	session := NewSession(SessionConfig{InitialStreamCredit: 2})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	mustEnqueue(t, session, 1, "one")
	mustEnqueue(t, session, 1, "two")
	if _, err := session.EnqueueData(1, []byte("blocked")); !errors.Is(err, ErrStreamCreditExhausted) {
		t.Fatalf("credit error = %v, want %v", err, ErrStreamCreditExhausted)
	}

	if err := session.Ack(1, 2); err != nil {
		t.Fatalf("ack: %v", err)
	}
	mustEnqueue(t, session, 1, "three")
	snapshot := session.Snapshot()
	if snapshot.FramesAcked != 2 || snapshot.Streams[1].Acked != 2 || snapshot.Streams[1].Credit != 1 {
		t.Fatalf("ack metrics = %+v stream=%+v", snapshot, snapshot.Streams[1])
	}
}

func TestSessionCreditIsCapped(t *testing.T) {
	session := NewSession(SessionConfig{InitialStreamCredit: 1, MaxStreamCredit: 2})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	if err := session.AddCredit(1, 100); err != nil {
		t.Fatalf("add credit: %v", err)
	}
	snapshot := session.Snapshot()
	if snapshot.Streams[1].Credit != 2 {
		t.Fatalf("credit = %d, want cap 2", snapshot.Streams[1].Credit)
	}
}

func TestSessionCloseRejectsNewData(t *testing.T) {
	session := NewSession(SessionConfig{})
	mustOpenStream(t, session, 1, TrafficClassReliable)
	session.Close()
	if _, err := session.EnqueueData(1, []byte("after-close")); !errors.Is(err, ErrSessionClosed) {
		t.Fatalf("enqueue after close err = %v, want %v", err, ErrSessionClosed)
	}
}

func mustOpenStream(t *testing.T, session *Session, streamID uint64, trafficClass TrafficClass) {
	t.Helper()
	if err := session.OpenStream(streamID, trafficClass); err != nil {
		t.Fatalf("open stream %d: %v", streamID, err)
	}
}

func mustEnqueue(t *testing.T, session *Session, streamID uint64, payload string) Frame {
	t.Helper()
	frame, err := session.EnqueueData(streamID, []byte(payload))
	if err != nil {
		t.Fatalf("enqueue stream %d: %v", streamID, err)
	}
	return frame
}