package chatclient import ( "context" "errors" "fmt" "io" "net" "strings" "sync" "testing" "time" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" "google.golang.org/grpc/test/bufconn" pb "gitea.fe.kaien.ru/kaien/l-ai-chat-proto" ) type fakeChatServer struct { pb.UnimplementedChatServiceServer mu sync.Mutex authByMethod map[string][]string postReqs []*pb.AgentPostMessageRequest genericPosts int ackCalls []ackCall ackErrors map[int64][]error ackFunc func(context.Context, *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) subscribeCount int subscribeFunc func(*pb.SubscribeRequest, pb.ChatService_SubscribeServer) error historyPages []*pb.HistoryResponse historyReqs []*pb.HistoryRequest } type ackCall struct { topic string subID string seq int64 } func newFake() *fakeChatServer { return &fakeChatServer{ authByMethod: make(map[string][]string), ackErrors: make(map[int64][]error), } } func (f *fakeChatServer) record(ctx context.Context, method string) { md, _ := metadata.FromIncomingContext(ctx) f.mu.Lock() defer f.mu.Unlock() f.authByMethod[method] = append([]string(nil), md.Get("authorization")...) } func (f *fakeChatServer) ListTopics(ctx context.Context, _ *pb.ListTopicsRequest) (*pb.ListTopicsResponse, error) { f.record(ctx, "ListTopics") return &pb.ListTopicsResponse{Topics: []*pb.Topic{{Id: "topic-1", Slug: "general"}}}, nil } func (f *fakeChatServer) GetTopic(ctx context.Context, _ *pb.GetTopicRequest) (*pb.Topic, error) { f.record(ctx, "GetTopic") return &pb.Topic{Id: "topic-1", Slug: "general"}, nil } func (f *fakeChatServer) GetMessage(ctx context.Context, req *pb.GetMessageRequest) (*pb.Message, error) { f.record(ctx, "GetMessage") return &pb.Message{Id: req.GetMessageId(), Sequence: 1}, nil } func (f *fakeChatServer) PostMessage(context.Context, *pb.PostMessageRequest) (*pb.PostMessageResponse, error) { f.mu.Lock() f.genericPosts++ f.mu.Unlock() return nil, status.Error(codes.Unimplemented, "generic post is fail-closed") } func (f *fakeChatServer) AgentPostMessage(ctx context.Context, req *pb.AgentPostMessageRequest) (*pb.PostMessageResponse, error) { f.record(ctx, "AgentPostMessage") f.mu.Lock() f.postReqs = append(f.postReqs, req) f.mu.Unlock() return &pb.PostMessageResponse{ Message: &pb.Message{Id: "posted", TopicId: req.GetTopicId(), Content: req.GetContent(), Sequence: 10}, }, nil } func (f *fakeChatServer) History(ctx context.Context, req *pb.HistoryRequest) (*pb.HistoryResponse, error) { f.record(ctx, "History") f.mu.Lock() defer f.mu.Unlock() f.historyReqs = append(f.historyReqs, req) if len(f.historyPages) == 0 { return &pb.HistoryResponse{}, nil } page := f.historyPages[0] f.historyPages = f.historyPages[1:] return page, nil } func (f *fakeChatServer) Subscribe(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { f.record(stream.Context(), "Subscribe") f.mu.Lock() f.subscribeCount++ fn := f.subscribeFunc f.mu.Unlock() if fn == nil { return status.Error(codes.Unavailable, "no subscribe script") } return fn(req, stream) } func (f *fakeChatServer) Acknowledge(ctx context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { f.record(ctx, "Acknowledge") f.mu.Lock() fn := f.ackFunc f.mu.Unlock() if fn != nil { return fn(ctx, req) } f.mu.Lock() f.ackCalls = append(f.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) errs := f.ackErrors[req.GetLastAckSequence()] if len(errs) > 0 { err := errs[0] f.ackErrors[req.GetLastAckSequence()] = errs[1:] f.mu.Unlock() return nil, err } f.mu.Unlock() return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil } func newClientForTest(t *testing.T, fake *fakeChatServer, opts ...Option) *Client { t.Helper() lis := bufconn.Listen(1 << 20) srv := grpc.NewServer() pb.RegisterChatServiceServer(srv, fake) go func() { _ = srv.Serve(lis) }() t.Cleanup(func() { srv.Stop() _ = lis.Close() }) all := []Option{ WithStaticToken("test-token"), WithInsecureTransport(), WithGRPCDialOptions(grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { return lis.DialContext(ctx) })), withJitterForTest(func(d time.Duration) time.Duration { return d }), withSleeperForTest(func(ctx context.Context, d time.Duration) error { select { case <-ctx.Done(): return ctx.Err() default: return nil } }), } all = append(all, opts...) c, err := New(context.Background(), "passthrough:///bufnet", all...) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = c.Close() }) return c } func msg(seq int64) *pb.DeliveryMessage { return &pb.DeliveryMessage{ Message: &pb.Message{Id: fmt.Sprintf("m-%d", seq), TopicId: "topic", Sequence: seq}, SubscriptionId: "sub-1", } } func recvDelivery(t *testing.T, sub *Subscription) Delivery { t.Helper() select { case d := <-sub.Deliveries(): return d case <-time.After(2 * time.Second): t.Fatal("timed out waiting for delivery") return Delivery{} } } func waitDone(t *testing.T, sub *Subscription) { t.Helper() select { case <-sub.Done(): case <-time.After(2 * time.Second): t.Fatal("timed out waiting for subscription done") } } func waitACKs(t *testing.T, fake *fakeChatServer, want ...int64) { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { fake.mu.Lock() got := make([]int64, len(fake.ackCalls)) for i, c := range fake.ackCalls { got[i] = c.seq } fake.mu.Unlock() if len(got) >= len(want) { ok := true for i := range want { if got[i] != want[i] { ok = false } } if ok { return } } time.Sleep(10 * time.Millisecond) } fake.mu.Lock() defer fake.mu.Unlock() t.Fatalf("ack calls = %+v, want prefix %v", fake.ackCalls, want) } func TestAuthMetadataUnarySubscribeAcknowledgeAndEmptyTokenFailsClosed(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { if err := stream.Send(&pb.DeliveryMessage{Ready: true}); err != nil { return err } if err := stream.Send(msg(100)); err != nil { return err } <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) if _, err := c.ListTopics(context.Background()); err != nil { t.Fatal(err) } sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } d := recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d); err != nil { t.Fatal(err) } waitACKs(t, fake, 100) fake.mu.Lock() defer fake.mu.Unlock() for _, method := range []string{"ListTopics", "Subscribe", "Acknowledge"} { got := fake.authByMethod[method] if len(got) != 1 || got[0] != "Bearer test-token" { t.Fatalf("%s auth metadata = %v", method, got) } } empty := newClientForTest(t, newFake(), WithStaticToken("")) _, err = empty.ListTopics(context.Background()) if status.Code(err) != codes.Unauthenticated { t.Fatalf("empty token: got %v", err) } if err != nil && strings.Contains(err.Error(), "test-token") { t.Fatalf("error leaked token: %v", err) } } func TestPostMessageUsesAgentPostMessageAndPreservesIdempotencyKey(t *testing.T) { fake := newFake() c := newClientForTest(t, fake) _, err := c.PostMessage(context.Background(), PostMessageRequest{ TopicID: "topic", Content: "payload", IdempotencyKey: "logical-key", }) if err != nil { t.Fatal(err) } fake.mu.Lock() defer fake.mu.Unlock() if fake.genericPosts != 0 { t.Fatal("generic PostMessage must not be called") } if len(fake.postReqs) != 1 || fake.postReqs[0].GetIdempotencyKey() != "logical-key" { t.Fatalf("AgentPostMessage reqs = %+v", fake.postReqs) } if fake.postReqs[0].GetDelegation() != nil { t.Fatalf("expected nil delegation, got %+v", fake.postReqs[0].GetDelegation()) } } func TestPostMessageForwardsDelegationMetadata(t *testing.T) { fake := newFake() c := newClientForTest(t, fake) delegation := &pb.AgentDelegation{ TargetAuthorId: "22222222-2222-2222-2222-222222222222", Kind: "INVESTIGATE", } _, err := c.PostMessage(context.Background(), PostMessageRequest{ TopicID: "topic", Content: "please investigate", IdempotencyKey: "deleg-key", Delegation: delegation, }) if err != nil { t.Fatal(err) } fake.mu.Lock() defer fake.mu.Unlock() if fake.genericPosts != 0 { t.Fatal("generic PostMessage must not be called") } if len(fake.postReqs) != 1 { t.Fatalf("AgentPostMessage reqs = %+v", fake.postReqs) } got := fake.postReqs[0].GetDelegation() if got == nil { t.Fatal("expected delegation to be forwarded") } if got.GetTargetAuthorId() != delegation.TargetAuthorId || got.GetKind() != delegation.Kind { t.Fatalf("delegation = %+v, want %+v", got, delegation) } } func TestOrderedCumulativeACKWithSequenceGaps(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) for _, seq := range []int64{100, 104, 111} { _ = stream.Send(msg(seq)) } <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } d100, d104, d111 := recvDelivery(t, sub), recvDelivery(t, sub), recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d104); err != nil { t.Fatal(err) } time.Sleep(30 * time.Millisecond) fake.mu.Lock() if len(fake.ackCalls) != 0 { t.Fatalf("ACK advanced over unprocessed 100: %+v", fake.ackCalls) } fake.mu.Unlock() if err := sub.MarkProcessed(context.Background(), d100); err != nil { t.Fatalf("mark 100: %v", err) } waitACKs(t, fake, 104) if err := sub.MarkProcessed(context.Background(), d111); err != nil { t.Fatalf("mark 111: %v", err) } waitACKs(t, fake, 104, 111) } func TestForgedDeliveryRejectedWithoutACK(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } err = sub.MarkProcessed(context.Background(), Delivery{ Message: &pb.Message{Sequence: 999999}, Replay: true, }) if !errors.Is(err, ErrUnknownDelivery) { t.Fatalf("forged delivery MarkProcessed = %v", err) } fake.mu.Lock() defer fake.mu.Unlock() if len(fake.ackCalls) != 0 { t.Fatalf("forged delivery produced ACKs: %+v", fake.ackCalls) } } func TestCrossSubscriptionDeliveryRejectedWithoutACK(t *testing.T) { fake := newFake() fake.subscribeFunc = func(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) dm := msg(100) dm.Message.TopicId = req.GetTopicId() dm.SubscriptionId = "sub-" + req.GetTopicId() _ = stream.Send(dm) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) subA, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "a"}) subB, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "b"}) dA := recvDelivery(t, subA) _ = recvDelivery(t, subB) if err := subB.MarkProcessed(context.Background(), dA); !errors.Is(err, ErrUnknownDelivery) { t.Fatalf("cross-subscription MarkProcessed = %v", err) } fake.mu.Lock() defer fake.mu.Unlock() if len(fake.ackCalls) != 0 { t.Fatalf("cross-subscription delivery produced ACKs: %+v", fake.ackCalls) } } func TestMutatedDeliveryMessageCannotAlterACKSequence(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) _ = stream.Send(msg(100)) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d := recvDelivery(t, sub) d.Message.Sequence = 111 if err := sub.MarkProcessed(context.Background(), d); err != nil { t.Fatal(err) } waitACKs(t, fake, 100) } func TestDuplicateMarkProcessedIsHarmless(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) _ = stream.Send(msg(100)) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d := recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d); err != nil { t.Fatal(err) } waitACKs(t, fake, 100) if err := sub.MarkProcessed(context.Background(), d); err != nil { t.Fatalf("duplicate MarkProcessed = %v", err) } fake.mu.Lock() defer fake.mu.Unlock() if len(fake.ackCalls) != 1 || fake.ackCalls[0].seq != 100 { t.Fatalf("duplicate MarkProcessed changed ACKs: %+v", fake.ackCalls) } } func TestLargerOutOfOrderCompletionDoesNotSkipEarliestPending(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) for _, seq := range []int64{100, 104, 111, 118, 130} { _ = stream.Send(msg(seq)) } <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) deliveries := map[int64]Delivery{} for i := 0; i < 5; i++ { d := recvDelivery(t, sub) deliveries[d.Message.GetSequence()] = d } for _, seq := range []int64{118, 104} { _ = sub.MarkProcessed(context.Background(), deliveries[seq]) } time.Sleep(30 * time.Millisecond) fake.mu.Lock() if len(fake.ackCalls) != 0 { t.Fatalf("ACK skipped earliest pending: %+v", fake.ackCalls) } fake.mu.Unlock() if err := sub.MarkProcessed(context.Background(), deliveries[100]); err != nil { t.Fatalf("mark 100: %v", err) } waitACKs(t, fake, 104) if err := sub.MarkProcessed(context.Background(), deliveries[130]); err != nil { t.Fatalf("mark 130: %v", err) } time.Sleep(30 * time.Millisecond) fake.mu.Lock() if len(fake.ackCalls) != 1 { t.Fatalf("ACK skipped 111: %+v", fake.ackCalls) } fake.mu.Unlock() if err := sub.MarkProcessed(context.Background(), deliveries[111]); err != nil { t.Fatalf("mark 111: %v", err) } waitACKs(t, fake, 104, 130) } func TestACKTransientFailureIsRetriedWithoutLosingBoundary(t *testing.T) { fake := newFake() fake.ackErrors[104] = []error{status.Error(codes.Unavailable, "try again")} fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) _ = stream.Send(msg(100)) _ = stream.Send(msg(104)) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d100); err != nil { t.Fatalf("mark 100: %v", err) } if err := sub.MarkProcessed(context.Background(), d104); err != nil { t.Fatalf("mark 104: %v", err) } waitACKs(t, fake, 104, 104) } func TestACKResponseLostAmbiguityReconnectUsesServerCursor(t *testing.T) { fake := newFake() dropFirst := make(chan struct{}) var dropOnce sync.Once var calls int fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { calls++ _ = stream.Send(&pb.DeliveryMessage{Ready: true}) if calls == 1 { _ = stream.Send(msg(100)) _ = stream.Send(msg(104)) <-dropFirst return status.Error(codes.Unavailable, "ack response was lost with the stream") } dm := msg(111) dm.SubscriptionId = "sub-2" _ = stream.Send(dm) <-stream.Context().Done() return stream.Context().Err() } fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { fake.mu.Lock() fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) fake.mu.Unlock() if req.GetLastAckSequence() == 104 { // Simulate "server applied ACK, client lost response"; the next // stream script starts after that durable cursor and does not // replay 100/104. dropOnce.Do(func() { close(dropFirst) }) return nil, status.Error(codes.Unavailable, "response lost") } return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil } c := newClientForTest(t, fake, WithReconnectBackoff(10*time.Millisecond, 10*time.Millisecond), withSleeperForTest(sleepContext), ) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d104); err != nil { t.Fatal(err) } if err := sub.MarkProcessed(context.Background(), d100); err != nil { t.Fatal(err) } d111 := recvDelivery(t, sub) if d111.Message.GetSequence() != 111 { t.Fatalf("reconnect must follow server cursor, got seq %d", d111.Message.GetSequence()) } if err := sub.MarkProcessed(context.Background(), d111); err != nil { t.Fatal(err) } waitACKs(t, fake, 104, 104, 111) } func TestACKPrefixGrowthWhileACKInFlightAdvancesAfterSuccess(t *testing.T) { fake := newFake() ack100Started := make(chan struct{}) release100 := make(chan struct{}) var once sync.Once fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) _ = stream.Send(msg(100)) _ = stream.Send(msg(104)) <-stream.Context().Done() return stream.Context().Err() } fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { fake.mu.Lock() fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) fake.mu.Unlock() if req.GetLastAckSequence() == 100 { once.Do(func() { close(ack100Started) }) <-release100 } return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d100); err != nil { t.Fatal(err) } <-ack100Started if err := sub.MarkProcessed(context.Background(), d104); err != nil { t.Fatal(err) } close(release100) waitACKs(t, fake, 100, 104) } func TestOldACKResultCannotCorruptReconnectedGeneration(t *testing.T) { fake := newFake() dropFirst := make(chan struct{}) oldACKStarted := make(chan struct{}) releaseOldACK := make(chan struct{}) var oldOnce sync.Once var calls int fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { calls++ _ = stream.Send(&pb.DeliveryMessage{Ready: true}) if calls == 1 { _ = stream.Send(msg(100)) <-dropFirst return status.Error(codes.Unavailable, "generation A failed") } dm := msg(100) dm.Replay = true dm.SubscriptionId = "sub-2" _ = stream.Send(dm) <-stream.Context().Done() return stream.Context().Err() } fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) { fake.mu.Lock() fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()}) fake.mu.Unlock() if req.GetSubscriptionId() == "sub-1" { oldOnce.Do(func() { close(oldACKStarted) }) <-releaseOldACK return nil, status.Error(codes.Aborted, "old subscription_id rejected") } return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil } c := newClientForTest(t, fake, WithReconnectBackoff(10*time.Millisecond, 10*time.Millisecond), withSleeperForTest(sleepContext), ) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) old100 := recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), old100); err != nil { t.Fatal(err) } <-oldACKStarted close(dropFirst) new100 := recvDelivery(t, sub) close(releaseOldACK) if err := sub.MarkProcessed(context.Background(), new100); err != nil { t.Fatal(err) } waitACKs(t, fake, 100, 100) if err := sub.Err(); err != nil { t.Fatalf("old ACK result corrupted generation B: %v", err) } fake.mu.Lock() defer fake.mu.Unlock() if fake.ackCalls[0].subID != "sub-1" || fake.ackCalls[1].subID != "sub-2" { t.Fatalf("ACKs used wrong subscription ids: %+v", fake.ackCalls) } } func TestReconnectReplayAndStaleGenerationFencing(t *testing.T) { fake := newFake() var calls int fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { calls++ _ = stream.Send(&pb.DeliveryMessage{Ready: true}) if calls == 1 { _ = stream.Send(msg(100)) _ = stream.Send(msg(104)) return status.Error(codes.Unavailable, "drop") } dm := msg(104) dm.Replay = true dm.SubscriptionId = "sub-2" _ = stream.Send(dm) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d100, old104 := recvDelivery(t, sub), recvDelivery(t, sub) if err := sub.MarkProcessed(context.Background(), d100); err != nil { t.Fatalf("mark 100: %v", err) } waitACKs(t, fake, 100) new104 := recvDelivery(t, sub) if !new104.Replay { t.Fatal("reconnected delivery must expose replay=true") } if err := sub.MarkProcessed(context.Background(), old104); !errors.Is(err, ErrStaleDelivery) { t.Fatalf("old generation MarkProcessed = %v", err) } if err := sub.MarkProcessed(context.Background(), new104); err != nil { t.Fatal(err) } waitACKs(t, fake, 100, 104) } func TestStreamReplacementAbortedIsTerminalNoReconnect(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) return status.Error(codes.Aborted, "subscription replaced") } c := newClientForTest(t, fake) sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } waitDone(t, sub) if !errors.Is(sub.Err(), ErrSubscriptionReplaced) { t.Fatalf("Err = %v", sub.Err()) } fake.mu.Lock() defer fake.mu.Unlock() if fake.subscribeCount != 1 { t.Fatalf("replacement triggered reconnect storm: %d subscribes", fake.subscribeCount) } } func TestReplacementBetweenTwoSDKSubscriptionsDoesNotPingPong(t *testing.T) { fake := newFake() replaceFirst := make(chan struct{}) secondReady := make(chan struct{}) fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { fake.mu.Lock() call := fake.subscribeCount fake.mu.Unlock() _ = stream.Send(&pb.DeliveryMessage{Ready: true}) if call == 1 { <-replaceFirst return status.Error(codes.Aborted, "subscription replaced") } close(secondReady) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) first, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } second, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } <-secondReady close(replaceFirst) waitDone(t, first) if !errors.Is(first.Err(), ErrSubscriptionReplaced) { t.Fatalf("first Err = %v", first.Err()) } if err := second.Close(); err != nil { t.Fatal(err) } fake.mu.Lock() defer fake.mu.Unlock() if fake.subscribeCount != 2 { t.Fatalf("replacement caused reconnect ping-pong: %d subscribes", fake.subscribeCount) } } func TestAuthTerminalErrorsDoNotEndlesslyReconnect(t *testing.T) { for _, code := range []codes.Code{codes.Unauthenticated, codes.PermissionDenied} { t.Run(code.String(), func(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, _ pb.ChatService_SubscribeServer) error { return status.Error(code, "terminal") } c := newClientForTest(t, fake) if _, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}); status.Code(err) != code { t.Fatalf("Subscribe error = %v", err) } fake.mu.Lock() defer fake.mu.Unlock() if fake.subscribeCount != 1 { t.Fatalf("terminal auth reconnects: %d", fake.subscribeCount) } }) } } func TestInitialSubscribeFailuresBeforeReadyReturnImmediately(t *testing.T) { cases := []struct { name string err error code codes.Code }{ {"Unavailable", status.Error(codes.Unavailable, "transient outage"), codes.Unavailable}, {"EOF", io.EOF, codes.Unknown}, {"DeadlineExceeded", status.Error(codes.DeadlineExceeded, "deadline"), codes.DeadlineExceeded}, {"Unauthenticated", status.Error(codes.Unauthenticated, "bad token"), codes.Unauthenticated}, {"PermissionDenied", status.Error(codes.PermissionDenied, "denied"), codes.PermissionDenied}, {"NotFound", status.Error(codes.NotFound, "missing topic"), codes.NotFound}, {"InvalidArgument", status.Error(codes.InvalidArgument, "bad topic"), codes.InvalidArgument}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, _ pb.ChatService_SubscribeServer) error { return tc.err } c := newClientForTest(t, fake) _, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if status.Code(err) != tc.code { t.Fatalf("Subscribe error = %v, code %v want %v", err, status.Code(err), tc.code) } fake.mu.Lock() defer fake.mu.Unlock() if fake.subscribeCount != 1 { t.Fatalf("initial failure retried before READY: %d subscribes", fake.subscribeCount) } }) } } func TestTransientReconnect(t *testing.T) { fake := newFake() var calls int fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { calls++ _ = stream.Send(&pb.DeliveryMessage{Ready: true}) if calls == 1 { return status.Error(codes.Unavailable, "transient") } _ = stream.Send(msg(200)) <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } d := recvDelivery(t, sub) if d.Message.GetSequence() != 200 { t.Fatalf("delivery seq = %d", d.Message.GetSequence()) } } func TestBoundedBufferShutdownDoesNotDropOrPanic(t *testing.T) { fake := newFake() fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) for _, seq := range []int64{1, 2, 3} { _ = stream.Send(msg(seq)) } <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake, WithSubscriptionBufferSize(1)) sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) if err != nil { t.Fatal(err) } if err := sub.Close(); err != nil { t.Fatal(err) } } func TestMultipleTopicSubscriptionsAreIndependent(t *testing.T) { fake := newFake() fake.subscribeFunc = func(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) dm := msg(1) dm.Message.TopicId = req.GetTopicId() dm.SubscriptionId = "sub-" + req.GetTopicId() _ = stream.Send(dm) if req.GetTopicId() == "general" { return status.Error(codes.Unavailable, "one topic failed") } <-stream.Context().Done() return stream.Context().Err() } c := newClientForTest(t, fake) general, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "general"}) security, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "security"}) if recvDelivery(t, general).Message.GetTopicId() != "general" { t.Fatal("general delivery mismatch") } ds := recvDelivery(t, security) if ds.Message.GetTopicId() != "security" { t.Fatal("security delivery mismatch") } if err := security.MarkProcessed(context.Background(), ds); err != nil { t.Fatal(err) } waitACKs(t, fake, 1) } func TestCloseRacesWithReconnectRecvAckAndMarkProcessed(t *testing.T) { fake := newFake() ackEntered := make(chan struct{}) releaseACK := make(chan struct{}) fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error { _ = stream.Send(&pb.DeliveryMessage{Ready: true}) _ = stream.Send(msg(1)) <-stream.Context().Done() return stream.Context().Err() } fake.ackErrors[1] = []error{status.Error(codes.Unavailable, "retry")} c := newClientForTest(t, fake, withSleeperForTest(func(ctx context.Context, d time.Duration) error { select { case <-ctx.Done(): return ctx.Err() default: return nil } })) origAck := fake.Acknowledge _ = origAck sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}) d := recvDelivery(t, sub) go func() { _ = sub.MarkProcessed(context.Background(), d) }() close(ackEntered) <-ackEntered close(releaseACK) go func() { _ = sub.MarkProcessed(context.Background(), d) }() if err := sub.Close(); err != nil { t.Fatal(err) } } func TestHistoryPagerUsesReturnedSequenceAcrossGaps(t *testing.T) { fake := newFake() fake.historyPages = []*pb.HistoryResponse{ {Messages: []*pb.Message{{Sequence: 100}, {Sequence: 104}, {Sequence: 111}}, HasMore: true}, {Messages: []*pb.Message{{Sequence: 118}}, HasMore: false}, } c := newClientForTest(t, fake) pager, err := c.NewHistoryPager(HistoryRequest{TopicID: "topic", AfterSequence: 99, Limit: 3}) if err != nil { t.Fatal(err) } first, err := pager.Next(context.Background()) if err != nil { t.Fatal(err) } if first.Messages[len(first.Messages)-1].GetSequence() != 111 { t.Fatal("first page setup broken") } second, err := pager.Next(context.Background()) if err != nil { t.Fatal(err) } if len(second.Messages) != 1 || second.Messages[0].GetSequence() != 118 { t.Fatalf("second page = %+v", second.Messages) } fake.mu.Lock() defer fake.mu.Unlock() if len(fake.historyReqs) != 2 || fake.historyReqs[1].GetAfterSequence() != 111 { t.Fatalf("history reqs = %+v", fake.historyReqs) } }