dingtalk-workspace-cli/internal/event/source/transient_auth_reconnect_test.go

// Copyright 2026 Alibaba Group
// Licensed under the Apache License, Version 2.0 (the "License");

package source

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	authpkg "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/auth"
	dwsevent "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/event"
	"github.com/gorilla/websocket"
	"github.com/open-dingtalk/dingtalk-stream-sdk-go/payload"
)

func transientRefreshError() error {
	return &authpkg.HTTPStatusError{StatusCode: http.StatusServiceUnavailable}
}

func terminalRefreshError() error {
	return &authpkg.HTTPStatusError{StatusCode: http.StatusUnauthorized}
}

func TestCrossPlatformCoveragePersonalRetryLogErrorReportsOnlySafeTransientStatus(t *testing.T) {
	cause := errors.New("personal source: ticket HTTP 401 secret detail")
	_, err := refreshRejectedSourceToken(context.Background(), func(context.Context, string) (string, error) {
		return "", transientRefreshError()
	}, "rejected", "personal source", cause)
	if got, want := personalRetryLogError(retryPersonal(err)), "personal source: token refresh HTTP 503"; got != want {
		t.Fatalf("personalRetryLogError() = %q, want %q", got, want)
	}
}

func TestCrossPlatformCoveragePersonalSourceRetriesTransientTokenResolutionFailure(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var calls atomic.Int32
	src, err := NewPersonal(PersonalConfig{
		AccessTokenProvider: func(context.Context) (string, error) {
			if calls.Add(1) == 2 {
				cancel()
			}
			return "", transientRefreshError()
		},
		ClientID:     "client",
		SourceID:     "open",
		TicketURL:    "https://ticket.invalid",
		ReconnectMin: time.Millisecond,
		ReconnectMax: time.Millisecond,
	})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(ctx, func(*dwsevent.RawEvent) {})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Start() error = %v, want context canceled after retry", err)
	}
	if calls.Load() != 2 || src.State().ReconnectCount != 1 {
		t.Fatalf("provider calls=%d reconnects=%d, want 2 calls and 1 reconnect", calls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePersonalSourceDoesNotRetryTerminalTokenResolutionFailure(t *testing.T) {
	var calls atomic.Int32
	src, err := NewPersonal(PersonalConfig{
		AccessTokenProvider: func(context.Context) (string, error) {
			calls.Add(1)
			return "", terminalRefreshError()
		},
		ClientID:  "client",
		SourceID:  "open",
		TicketURL: "https://ticket.invalid",
	})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(context.Background(), func(*dwsevent.RawEvent) {})
	if authpkg.ClassifyRefreshFailure(err) != authpkg.RefreshFailureTerminal {
		t.Fatalf("Start() error = %v, want terminal refresh failure", err)
	}
	if calls.Load() != 1 || src.State().ReconnectCount != 0 {
		t.Fatalf("provider calls=%d reconnects=%d, want 1 call and no reconnect", calls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePersonalSourceRetriesTransientRejectedTokenRefresh(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var ticketCalls atomic.Int32
	var refreshCalls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		ticketCalls.Add(1)
		w.WriteHeader(http.StatusUnauthorized)
	}))
	defer srv.Close()

	src, err := NewPersonal(PersonalConfig{
		AccessTokenProvider: func(context.Context) (string, error) { return "old-token", nil },
		ForceRefreshToken: func(_ context.Context, rejected string) (string, error) {
			if rejected != "old-token" {
				t.Fatalf("rejected token = %q, want old-token", rejected)
			}
			if refreshCalls.Add(1) == 2 {
				cancel()
			}
			return "", transientRefreshError()
		},
		ClientID:     "client",
		SourceID:     "open",
		TicketURL:    srv.URL,
		HTTPClient:   srv.Client(),
		ReconnectMin: time.Millisecond,
		ReconnectMax: time.Millisecond,
	})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(ctx, func(*dwsevent.RawEvent) {})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Start() error = %v, want context canceled after retry", err)
	}
	if ticketCalls.Load() != 2 || refreshCalls.Load() != 2 || src.State().ReconnectCount != 1 {
		t.Fatalf("ticket calls=%d refresh calls=%d reconnects=%d", ticketCalls.Load(), refreshCalls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePersonalRetryLogErrorFallsBackToNetworkMessage(t *testing.T) {
	err := retryPersonal(fmt.Errorf("personal source: resolve access token: %w", errors.New("dial tcp: lookup oauth.invalid")))
	if got, want := personalRetryLogError(err), "personal source: token refresh: temporary network error"; got != want {
		t.Fatalf("personalRetryLogError() = %q, want %q", got, want)
	}
}

func TestCrossPlatformCoveragePortalStageErrorNilAndUnwrap(t *testing.T) {
	var nilErr *portalStageError
	if got, want := nilErr.Error(), "source: portal stream failed"; got != want {
		t.Fatalf("nil stage error = %q, want %q", got, want)
	}
	if nilErr.Unwrap() != nil {
		t.Fatal("nil stage error should unwrap to nil")
	}
	cause := errors.New("cause")
	stageErr := &portalStageError{stage: "stream_read", retryable: true, cause: cause}
	if !errors.Is(stageErr, cause) {
		t.Fatalf("stage error should unwrap to cause: %v", stageErr)
	}
}

func TestCrossPlatformCoveragePortalSourceRetriesTransientTokenResolutionFailure(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var calls atomic.Int32
	src, err := New(Config{PortalTicket: &PortalTicketConfig{
		TicketURL: "https://ticket.invalid",
		AccessTokenProvider: func(context.Context) (string, error) {
			if calls.Add(1) == 2 {
				cancel()
			}
			return "", transientRefreshError()
		},
		SourceID: "open",
		// Min above max exercises the reconnect clamp.
		ReconnectMin: 2 * time.Millisecond,
		ReconnectMax: time.Millisecond,
	}})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(ctx, func(*dwsevent.RawEvent) {})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Start() error = %v, want context canceled after retry", err)
	}
	if calls.Load() != 2 || src.State().ReconnectCount != 1 {
		t.Fatalf("provider calls=%d reconnects=%d, want 2 calls and 1 reconnect", calls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePortalSourceResetsBackoffAfterAckedAttempt(t *testing.T) {
	var ticketCalls atomic.Int32
	upgrader := websocket.Upgrader{}
	var wsURL string
	mux := http.NewServeMux()
	mux.HandleFunc("/ticket", func(w http.ResponseWriter, _ *http.Request) {
		if ticketCalls.Add(1) > 1 {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		_, _ = io.WriteString(w, `{"endpoint":`+strconvQuote(wsURL)+`,"ticket":"t"}`)
	})
	mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		df := payload.DataFrame{Type: "event", Headers: payload.DataFrameHeader{payload.DataFrameHeaderKMessageId: "m"}, Data: `{}`}
		_ = conn.WriteJSON(df)
		// Wait for the ACK, then close so the read fails retryably with an
		// acked attempt behind it, which resets the reconnect backoff.
		_, _, _ = conn.ReadMessage()
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	wsURL = "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws"

	src, err := New(Config{PortalTicket: &PortalTicketConfig{
		TicketURL:    srv.URL + "/ticket",
		AccessToken:  "t",
		SourceID:     "open",
		HTTPClient:   srv.Client(),
		ReconnectMin: time.Millisecond,
		ReconnectMax: time.Millisecond,
	}})
	if err != nil {
		t.Fatal(err)
	}
	var events atomic.Int32
	err = src.Start(context.Background(), func(*dwsevent.RawEvent) { events.Add(1) })
	if err == nil || !strings.Contains(err.Error(), "portal ticket HTTP 400") {
		t.Fatalf("Start() error = %v, want fatal ticket HTTP 400 after reconnect", err)
	}
	if events.Load() != 1 || ticketCalls.Load() != 2 || src.State().ReconnectCount != 1 {
		t.Fatalf("events=%d ticket calls=%d reconnects=%d, want 1/2/1", events.Load(), ticketCalls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePortalSourceDoesNotRetryTerminalTokenResolutionFailure(t *testing.T) {
	var calls atomic.Int32
	src, err := New(Config{PortalTicket: &PortalTicketConfig{
		TicketURL: "https://ticket.invalid",
		AccessTokenProvider: func(context.Context) (string, error) {
			calls.Add(1)
			return "", terminalRefreshError()
		},
		SourceID: "open",
	}})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(context.Background(), func(*dwsevent.RawEvent) {})
	if authpkg.ClassifyRefreshFailure(err) != authpkg.RefreshFailureTerminal {
		t.Fatalf("Start() error = %v, want terminal refresh failure", err)
	}
	if calls.Load() != 1 || src.State().ReconnectCount != 0 {
		t.Fatalf("provider calls=%d reconnects=%d, want 1 call and no reconnect", calls.Load(), src.State().ReconnectCount)
	}
}

func TestCrossPlatformCoveragePortalSourceRetriesTransientRejectedTokenRefresh(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var ticketCalls atomic.Int32
	var refreshCalls atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		ticketCalls.Add(1)
		w.WriteHeader(http.StatusUnauthorized)
	}))
	defer srv.Close()

	src, err := New(Config{PortalTicket: &PortalTicketConfig{
		TicketURL:           srv.URL,
		AccessTokenProvider: func(context.Context) (string, error) { return "old-token", nil },
		ForceRefreshToken: func(_ context.Context, rejected string) (string, error) {
			if rejected != "old-token" {
				t.Fatalf("rejected token = %q, want old-token", rejected)
			}
			if refreshCalls.Add(1) == 2 {
				cancel()
			}
			return "", transientRefreshError()
		},
		SourceID:     "open",
		HTTPClient:   srv.Client(),
		ReconnectMin: time.Millisecond,
		ReconnectMax: time.Millisecond,
	}})
	if err != nil {
		t.Fatal(err)
	}
	err = src.Start(ctx, func(*dwsevent.RawEvent) {})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Start() error = %v, want context canceled after retry", err)
	}
	if ticketCalls.Load() != 2 || refreshCalls.Load() != 2 || src.State().ReconnectCount != 1 {
		t.Fatalf("ticket calls=%d refresh calls=%d reconnects=%d", ticketCalls.Load(), refreshCalls.Load(), src.State().ReconnectCount)
	}
}