CamTalk/backend/internal/ai/stt/deepgram_test.go

package stt

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"go.uber.org/zap"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// newMockDeepgram 创建模拟 Deepgram WebSocket 服务。
// 返回 httptest.Server 和对应的 ws:// URL。
func newMockDeepgram(t *testing.T, handler func(conn *websocket.Conn)) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Logf("upgrade error: %v", err)
			return
		}
		defer conn.Close()
		handler(conn)
	}))
	return srv
}

// wsToWss 将 http:// 转换为 ws://。
func wsToWss(httpURL string) string {
	return "ws" + strings.TrimPrefix(httpURL, "http")
}

func TestDeepgramService_Recognize_Success(t *testing.T) {
	srv := newMockDeepgram(t, func(conn *websocket.Conn) {
		// 读取音频数据
		_, _, err := conn.ReadMessage()
		if err != nil {
			t.Errorf("read audio: %v", err)
			return
		}

		// 发送中间结果（非 final）
		intermediate := deepgramResponse{
			IsFinal: false,
		}
		intermediate.Channel.Alternatives = []struct {
			Transcript string  `json:"transcript"`
			Confidence float64 `json:"confidence"`
		}{{Transcript: "你好", Confidence: 0.9}}
		data, _ := json.Marshal(intermediate)
		_ = conn.WriteMessage(websocket.TextMessage, data)

		// 发送最终结果
		final := deepgramResponse{
			IsFinal: true,
		}
		final.Channel.Alternatives = []struct {
			Transcript string  `json:"transcript"`
			Confidence float64 `json:"confidence"`
		}{{Transcript: "你好世界", Confidence: 0.95}}
		data, _ = json.Marshal(final)
		_ = conn.WriteMessage(websocket.TextMessage, data)

		// 等待客户端关闭
		_, _, _ = conn.ReadMessage()
	})
	defer srv.Close()

	svc := NewDeepgramService("test-key", "", wsToWss(srv.URL)+"/v1/listen", 0, zap.NewNop().Sugar())

	text, err := svc.Recognize(context.Background(), []byte("fake-pcm-audio"), Options{
		Encoding:   "pcm_s16le",
		SampleRate: 16000,
		Language:   "zh-CN",
	})
	if err != nil {
		t.Fatalf("Recognize() error: %v", err)
	}
	if text != "你好世界" {
		t.Errorf("Recognize() = %q, want %q", text, "你好世界")
	}
}

func TestDeepgramService_Recognize_EmptyAudio(t *testing.T) {
	svc := NewDeepgramService("test-key", "", "ws://localhost", 0, zap.NewNop().Sugar())
	_, err := svc.Recognize(context.Background(), nil, Options{})
	if err == nil {
		t.Fatal("Recognize() with empty audio should return error")
	}
}

func TestDeepgramService_Recognize_ConnectError(t *testing.T) {
	svc := NewDeepgramService("test-key", "", "ws://localhost:1", 0, zap.NewNop().Sugar())
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	_, err := svc.Recognize(ctx, []byte("audio"), Options{})
	if err == nil {
		t.Fatal("Recognize() with bad endpoint should return error")
	}
}

func TestDeepgramService_Recognize_Timeout(t *testing.T) {
	// 模拟一个永不响应的服务端
	srv := newMockDeepgram(t, func(conn *websocket.Conn) {
		// 读取音频但不发送任何结果，让客户端超时
		_, _, _ = conn.ReadMessage()
		time.Sleep(10 * time.Second)
	})
	defer srv.Close()

	svc := NewDeepgramService("test-key", "", wsToWss(srv.URL)+"/v1/listen", 0, zap.NewNop().Sugar())

	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	_, err := svc.Recognize(ctx, []byte("audio"), Options{})
	if err == nil {
		t.Fatal("Recognize() should timeout")
	}
}

func TestDeepgramService_Recognize_MultipleFinals(t *testing.T) {
	srv := newMockDeepgram(t, func(conn *websocket.Conn) {
		_, _, _ = conn.ReadMessage()

		// 发送多个 final 结果（多句话场景）
		for _, text := range []string{"你好", "世界"} {
			resp := deepgramResponse{IsFinal: true}
			resp.Channel.Alternatives = []struct {
				Transcript string  `json:"transcript"`
				Confidence float64 `json:"confidence"`
			}{{Transcript: text, Confidence: 0.9}}
			data, _ := json.Marshal(resp)
			_ = conn.WriteMessage(websocket.TextMessage, data)
		}

		_, _, _ = conn.ReadMessage()
	})
	defer srv.Close()

	svc := NewDeepgramService("test-key", "", wsToWss(srv.URL)+"/v1/listen", 0, zap.NewNop().Sugar())

	text, err := svc.Recognize(context.Background(), []byte("audio"), Options{})
	if err != nil {
		t.Fatalf("Recognize() error: %v", err)
	}
	if text != "你好世界" {
		t.Errorf("Recognize() = %q, want %q", text, "你好世界")
	}
}

func TestDeepgramService_buildURL(t *testing.T) {
	svc := NewDeepgramService("key", "", "wss://api.deepgram.com/v1/listen", 0, zap.NewNop().Sugar())

	tests := []struct {
		name string
		opts Options
		want []string // URL 中应包含的参数
	}{
		{
			name: "defaults",
			opts: Options{},
			want: []string{"encoding=pcm_s16le", "sample_rate=16000", "language=zh-CN"},
		},
		{
			name: "custom",
			opts: Options{Encoding: "wav", SampleRate: 44100, Language: "en"},
			want: []string{"encoding=wav", "sample_rate=44100", "language=en"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			u := svc.buildURL(tt.opts)
			for _, param := range tt.want {
				if !strings.Contains(u, param) {
					t.Errorf("buildURL() = %q, should contain %q", u, param)
				}
			}
		})
	}
}