larksuite--cli/internal/event/consume/startup_probe_test.go

// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT

package consume

import (
	"bufio"
	"bytes"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/larksuite/cli/internal/event/protocol"
)

type probeMockTransport struct {
	mu       sync.Mutex
	listener net.Listener
	addr     string

	wg    sync.WaitGroup
	conns []net.Conn
}

func newProbeMockTransport(t *testing.T) *probeMockTransport {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	return &probeMockTransport{listener: ln, addr: ln.Addr().String()}
}

func (m *probeMockTransport) Listen(addr string) (net.Listener, error) {
	return m.listener, nil
}

func (m *probeMockTransport) Dial(addr string) (net.Conn, error) {
	return net.Dial("tcp", m.addr)
}

func (m *probeMockTransport) Address(appID string) string { return m.addr }
func (m *probeMockTransport) Cleanup(addr string)         {}

func (m *probeMockTransport) trackConn(c net.Conn) {
	m.mu.Lock()
	m.conns = append(m.conns, c)
	m.mu.Unlock()
}

func (m *probeMockTransport) stop() {
	m.mu.Lock()
	_ = m.listener.Close()
	conns := append([]net.Conn(nil), m.conns...)
	m.conns = nil
	m.mu.Unlock()
	for _, c := range conns {
		_ = c.Close()
	}
	m.wg.Wait()
}

func runHealthyBus(t *testing.T, m *probeMockTransport) {
	t.Helper()
	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		probeConn, err := m.listener.Accept()
		if err != nil {
			return
		}
		m.trackConn(probeConn)
		br := bufio.NewReader(probeConn)
		line, _ := br.ReadBytes('\n')
		msg, _ := protocol.Decode(bytes.TrimRight(line, "\n"))
		if _, ok := msg.(*protocol.StatusQuery); ok {
			_ = protocol.Encode(probeConn, protocol.NewStatusResponse(12345, 10, 0, nil))
		}
		_ = probeConn.Close()

		realConn, err := m.listener.Accept()
		if err != nil {
			return
		}
		m.trackConn(realConn)
	}()
}

func runDeadBus(t *testing.T, m *probeMockTransport) {
	t.Helper()
	m.wg.Add(1)
	go func() {
		defer m.wg.Done()
		for {
			conn, err := m.listener.Accept()
			if err != nil {
				return
			}
			m.trackConn(conn)
			m.wg.Add(1)
			go func(c net.Conn) {
				defer m.wg.Done()
				buf := make([]byte, 4096)
				for {
					if _, err := c.Read(buf); err != nil {
						return
					}
				}
			}(conn)
		}
	}()
}

func TestProbeAndDialBusHealthy(t *testing.T) {
	m := newProbeMockTransport(t)
	t.Cleanup(m.stop)
	runHealthyBus(t, m)

	conn, err := probeAndDialBus(m, m.addr)
	if err != nil {
		t.Fatalf("expected success, got error: %v", err)
	}
	if conn == nil {
		t.Fatal("expected non-nil conn")
	}
	conn.Close()
}

func TestProbeAndDialBusUnresponsive(t *testing.T) {
	m := newProbeMockTransport(t)
	t.Cleanup(m.stop)
	runDeadBus(t, m)

	start := time.Now()
	conn, err := probeAndDialBus(m, m.addr)
	elapsed := time.Since(start)

	if err == nil {
		conn.Close()
		t.Fatal("expected error on unresponsive bus")
	}
	if elapsed > 3*time.Second {
		t.Errorf("expected ~2s timeout, got %v", elapsed)
	}
}