internal/home/client_cancellation_test.go

Rama main · réplica de lectura del repositorio original.

package home

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/router-for-me/CLIProxyAPI/v8/internal/config"
)

func TestManagedSubscriberCancellationPreservesFinalRelease(t *testing.T) {
	for _, existingReleaseClient := range []bool{false, true} {
		t.Run(fmt.Sprintf("existingReleaseClient=%v", existingReleaseClient), func(t *testing.T) {
			client, commands := newRedisCommandTestClient(t, func(args []string) string {
				switch strings.ToUpper(args[0]) {
				case "GET":
					return "$2\r\n{}\r\n"
				case "SUBSCRIBE":
					return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"
				case "LPUSH":
					return ":1\r\n"
				default:
					return balanceTestRedisReply(args)
				}
			})
			client.SetManagedLifetime(true)
			client.testOperationTimeout = time.Second
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1}
			var errInitial error
			ready := false
			errRun := client.RunConfigSubscriberLifetime(ctx, func([]byte) error {
				return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1})
			}, func() {
				ready = client.HeartbeatOK()
				if existingReleaseClient {
					errInitial = client.PushConcurrencyRelease(ctx, frame)
				}
				cancel()
			})
			if !ready || !errors.Is(errRun, context.Canceled) || errInitial != nil {
				t.Fatalf("subscriber cancellation: ready=%v error=%v initial release=%v", ready, errRun, errInitial)
			}
			if client.HeartbeatOK() || client.AmbiguousDispatch() {
				t.Fatal("graceful cancellation retained a heartbeat or introduced ambiguity")
			}
			if _, errDispatch := client.RPopAuth(context.Background(), "late-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) {
				t.Fatalf("graceful cancellation reopened dispatch: %v", errDispatch)
			}
			frame.ReleaseSeq = 2
			if errRelease := client.PushConcurrencyRelease(context.WithoutCancel(ctx), frame); errRelease != nil {
				t.Fatalf("final release after normal cancellation: %v", errRelease)
			}
			var releases []ConcurrencyReleaseFrame
			for _, args := range commands.All() {
				if len(args) == 3 && strings.EqualFold(args[0], "LPUSH") && args[1] == redisKeyConcurrencyRelease {
					var sent ConcurrencyReleaseFrame
					if errJSON := json.Unmarshal([]byte(args[2]), &sent); errJSON != nil {
						t.Fatal(errJSON)
					}
					releases = append(releases, sent)
				}
			}
			wantCount := 1
			if existingReleaseClient {
				wantCount++
			}
			if len(releases) != wantCount || releases[len(releases)-1] != frame {
				t.Fatalf("release frames = %#v", releases)
			}
			client.Close()
			if errRelease := client.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrDispatchFenced) {
				t.Fatalf("release after Close = %v, want fenced", errRelease)
			}
		})
	}
}

func TestManagedSubscriberReceiveFailureKeepsReleaseFenced(t *testing.T) {
	client, commands := newRedisCommandTestClient(t, func(args []string) string {
		switch strings.ToUpper(args[0]) {
		case "GET":
			return "$2\r\n{}\r\n"
		case "SUBSCRIBE":
			// A malformed RESP frame fails reception after the membership ACK.
			return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n?invalid\r\n"
		default:
			return balanceTestRedisReply(args)
		}
	})
	client.SetManagedLifetime(true)
	client.testOperationTimeout = time.Second
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	errRun := client.RunConfigSubscriberLifetime(ctx, func([]byte) error {
		return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1})
	}, nil)
	if errRun == nil || errors.Is(errRun, context.Canceled) {
		t.Fatalf("receive failure = %v", errRun)
	}
	// Later graceful cancellation must not undo a fence installed by a receive failure.
	cancel()
	if client.HeartbeatOK() || !client.dispatchFenced.Load() {
		t.Fatal("abnormal reception did not fence the lifetime")
	}
	frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1}
	if errRelease := client.PushConcurrencyRelease(context.WithoutCancel(ctx), frame); !errors.Is(errRelease, ErrDispatchFenced) {
		t.Fatalf("release after receive failure = %v, want fenced", errRelease)
	}
	if _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) {
		t.Fatalf("dispatch after receive failure = %v, want fenced", errDispatch)
	}
	if commands.CountCommandKey("LPUSH", redisKeyConcurrencyRelease) != 0 {
		t.Fatal("fenced lifetime sent a release")
	}
}