talk/internal/provider/edge_test.go

package provider

import (
	"context"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

// TestEdgeAvailable: Edge needs only an endpoint, no key. The constructor
// bakes in the default endpoint, so an Edge built via NewEdge is always
// available; only the zero value (built directly, without an endpoint) is
// not.
func TestEdgeAvailable(t *testing.T) {
	if (&Edge{}).Available(context.Background()) {
		t.Error("Available without endpoint (zero value) = true, want false")
	}
	p := NewEdge(ChainConfig{Getenv: envMap(map[string]string{
		EnvEdgeEndpoint: "wss://edge.example.com",
	})})
	if !p.Available(context.Background()) {
		t.Error("Available with endpoint = false, want true")
	}
	if !NewEdge(ChainConfig{Getenv: envMap(nil)}).Available(context.Background()) {
		t.Error("Available with default endpoint = false, want true (edge is the key-free fallback)")
	}
}

// TestEdgeConnectionFails (error-path mock): a guaranteed dead loopback
// port yields a clear error, immediately and without network.
func TestEdgeConnectionFails(t *testing.T) {
	p := NewEdge(ChainConfig{Getenv: envMap(map[string]string{
		EnvEdgeEndpoint: "ws://127.0.0.1:1/",
	})})
	_, err := p.Synthesize(context.Background(), "Hello", "de")
	if err == nil {
		t.Fatal("Synthesize against dead port: no error")
	}
	if !strings.Contains(err.Error(), "edge") || !strings.Contains(err.Error(), "WebSocket") {
		t.Errorf("error = %v, want a clear edge/WebSocket error", err)
	}
}

// TestEdgeInvalidScheme: a non-ws(s) endpoint is a clear error.
func TestEdgeInvalidScheme(t *testing.T) {
	p := NewEdge(ChainConfig{Getenv: envMap(map[string]string{
		EnvEdgeEndpoint: "http://127.0.0.1:1234/",
	})})
	_, err := p.Synthesize(context.Background(), "Hello", "de")
	if err == nil || !strings.Contains(err.Error(), "ws://") {
		t.Errorf("error = %v, want scheme hint ws://", err)
	}
}

// TestEdgeSSML: the SSML request contains the voice and the text stripped
// of control characters and XML-escaped; xml:lang is fixed to "en-US" as in
// the reference.
func TestEdgeSSML(t *testing.T) {
	got := edgeSSML("de-DE-KatjaNeural", "AT&T <customer> „done“")
	for _, want := range []string{
		"voice name='de-DE-KatjaNeural'",
		"xml:lang='en-US'",
		"AT&amp;T &lt;customer&gt;",
	} {
		if !strings.Contains(got, want) {
			t.Errorf("SSML = %q, want %q", got, want)
		}
	}
	if strings.Contains(got, "AT&T") {
		t.Errorf("SSML contains unescaped &: %q", got)
	}
}

// TestEdgeEscape: control characters (e.g. vertical tab from OCR texts) are
// replaced by spaces, & < > are escaped.
func TestEdgeEscape(t *testing.T) {
	in := "a\x0bb&c<d>\x08e"
	if want := "a b&amp;c&lt;d&gt; e"; edgeEscape(in) != want {
		t.Errorf("edgeEscape(%q) = %q, want %q", in, edgeEscape(in), want)
	}
}

// TestEdgeSpeechConfigMsg: output format (mono mp3) and path are correct.
func TestEdgeSpeechConfigMsg(t *testing.T) {
	msg := edgeSpeechConfigMsg(time.Unix(0, 0).UTC())
	if !strings.Contains(msg, "Path:speech.config") {
		t.Errorf("speech.config message without Path: %q", msg)
	}
	if !strings.Contains(msg, edgeOutputFormat) {
		t.Errorf("speech.config message without output format %q: %q", edgeOutputFormat, msg)
	}
}

// TestEdgeRequestURL: query parameters (token, Sec-MS-GEC, version) are
// added; the token does not end up in the redacted URL output.
func TestEdgeRequestURL(t *testing.T) {
	p := &Edge{Endpoint: "wss://edge.example.com/path"}
	raw, err := p.requestURL(time.Unix(1700000000, 0).UTC())
	if err != nil {
		t.Fatalf("requestURL: %v", err)
	}
	for _, want := range []string{
		"TrustedClientToken=" + edgeTrustedClientToken,
		"Sec-MS-GEC=",
		"Sec-MS-GEC-Version=" + edgeSecMSGECVersion,
	} {
		if !strings.Contains(raw, want) {
			t.Errorf("requestURL = %q, want %q", raw, want)
		}
	}
	if got := redactURL(raw); strings.Contains(got, edgeTrustedClientToken) {
		t.Errorf("redactURL leaks the token: %q", got)
	}
}

// TestEdgeSecMSGEC: format "<16 hex>-<64 hex>" (upper case), 5-minute grid
// (stable within the window, different across the boundary) and exact
// reference vectors of the edge-tts DRM implementation (v2).
func TestEdgeSecMSGEC(t *testing.T) {
	// Reference vectors: (unix seconds, expected token), generated with the
	// Python reference (edge-tts src/edge_tts/drm.py, as of 2026-10).
	vectors := []struct {
		ts   int64
		want string
	}{
		{1700000400, "73C7C85FE76CF347-452672184D494A9F80E179C2FBFEE2E98480C4D4AAFC1F4A577E43B023BEEEC9"},
		{1700000599, "73C7C85FE76CF347-452672184D494A9F80E179C2FBFEE2E98480C4D4AAFC1F4A577E43B023BEEEC9"},
		{1765800000, "B9B5447115B4FD1A-6BB36279F5B3D1B53F84F786AF05C81A86909ABE30E695496BBE88FFAE2E3E83"},
		{0, "52BC7A17B50DFEC9-28C4394775C8751CD2F6859C030C06DB1489ACAB3B3A7651831FC26895177ED6"},
	}
	for _, v := range vectors {
		if got := edgeSecMSGEC(time.Unix(v.ts, 0).UTC()); got != v.want {
			t.Errorf("edgeSecMSGEC(%d) =\n  %s\nwant\n  %s", v.ts, got, v.want)
		}
	}

	// Window stability (base at window start, mod 300 == 0).
	base := time.Unix(1700000400, 0).UTC()
	token := edgeSecMSGEC(base)
	if edgeSecMSGEC(base.Add(4*time.Minute)) != token {
		t.Error("Sec-MS-GEC differs within the 5-minute window")
	}
	if edgeSecMSGEC(base.Add(6*time.Minute)) == token {
		t.Error("Sec-MS-GEC does not change across the window boundary")
	}
	if len(token) != 16+1+64 || !strings.Contains(token, "-") {
		t.Errorf("token format = %q, want <16 hex>-<64 hex>", token)
	}
}

// TestEdgeRequestID: 32 hex characters, unique.
func TestEdgeRequestID(t *testing.T) {
	a, b := edgeRequestID(), edgeRequestID()
	if len(a) != 32 {
		t.Errorf("request ID length = %d, want 32", len(a))
	}
	if _, err := hex.DecodeString(a); err != nil {
		t.Errorf("request ID is not hex: %v", err)
	}
	if a == b {
		t.Error("request ID not unique")
	}
}

// TestEdgeSplitBinary: binary frame: 2 bytes big-endian header length, then
// header (Path:audio), then payload.
func TestEdgeSplitBinary(t *testing.T) {
	msg := []byte{0, 12, 'P', 'a', 't', 'h', ':', 'a', 'u', 'd', 'i', 'o', '\r', '\n', 1, 2, 3}
	header, payload, ok := splitEdgeBinary(msg)
	if !ok || header != "Path:audio\r\n" || string(payload) != "\x01\x02\x03" {
		t.Errorf("splitEdgeBinary = %q, %x, %v", header, payload, ok)
	}
	if _, _, ok := splitEdgeBinary([]byte{0, 99, 'x'}); ok {
		t.Error("oversized header length should be rejected")
	}
}

// fakeEdgeServer is a loopback fake of the Edge WSS protocol (no real
// network call; everything stays on 127.0.0.1). It checks the two client
// messages and answers per protocol: a binary audio frame with Path:audio,
// then turn.end. audio is the payload to deliver; closing closes the
// connection before turn.end (error path).
func fakeEdgeServer(t *testing.T, audio []byte, closing bool) (*httptest.Server, *fakeEdgeSeen) {
	t.Helper()
	seen := &fakeEdgeSeen{}
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if got := r.URL.Query().Get("TrustedClientToken"); got != edgeTrustedClientToken {
			t.Errorf("TrustedClientToken = %q", got)
		}
		if r.Header.Get("Origin") != edgeOrigin {
			t.Errorf("Origin = %q", r.Header.Get("Origin"))
		}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Errorf("Upgrade: %v", err)
			return
		}
		defer conn.Close()
		if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
			return
		}
		if _, cfg, err := conn.ReadMessage(); err != nil {
			t.Errorf("reading speech.config: %v", err)
		} else {
			seen.mu.Lock()
			seen.config = string(cfg)
			seen.mu.Unlock()
		}
		if _, req, err := conn.ReadMessage(); err != nil {
			t.Errorf("reading ssml: %v", err)
		} else {
			seen.mu.Lock()
			seen.ssml = string(req)
			seen.mu.Unlock()
		}
		if closing {
			return // close the connection without answering
		}
		header := "X-RequestId:dummy\r\nPath:audio\r\n"
		frame := make([]byte, 0, 2+len(header)+len(audio))
		frame = append(frame, byte(len(header)>>8), byte(len(header)))
		frame = append(frame, header...)
		frame = append(frame, audio...)
		if err := conn.WriteMessage(websocket.BinaryMessage, frame); err != nil {
			return
		}
		if err := conn.WriteMessage(websocket.TextMessage, []byte("X-RequestId:dummy\r\nPath:turn.end\r\n\r\n")); err != nil {
			return
		}
	}))
	t.Cleanup(ts.Close)
	return ts, seen
}

type fakeEdgeSeen struct {
	mu     sync.Mutex
	config string
	ssml   string
}

func (s *fakeEdgeSeen) snapshot() (configMsg, ssml string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.config, s.ssml
}

// TestEdgeLoopbackFakeServer: complete protocol run against the loopback
// fake (tested fakeably: handshake parameters, both client messages,
// binary frame parsing, turn.end termination).
func TestEdgeLoopbackFakeServer(t *testing.T) {
	wantAudio := []byte("\xff\xfb\x90\x00" + "edge-audio")
	ts, seen := fakeEdgeServer(t, wantAudio, false)

	p := NewEdge(ChainConfig{Getenv: envMap(map[string]string{
		EnvEdgeEndpoint: "ws://" + ts.Listener.Addr().String(),
	})})
	data, err := p.Synthesize(context.Background(), "Hello & done", "de")
	if err != nil {
		t.Fatalf("Synthesize against loopback fake: %v", err)
	}
	if string(data) != string(wantAudio) {
		t.Errorf("Synthesize = %x, want %x", data, wantAudio)
	}
	configMsg, ssml := seen.snapshot()
	if !strings.Contains(configMsg, "Path:speech.config") {
		t.Errorf("speech.config = %q", configMsg)
	}
	if !strings.Contains(ssml, "Path:ssml") {
		t.Errorf("ssml request without Path:ssml: %q", ssml)
	}
	if !strings.Contains(ssml, "voice name='de-DE-KatjaNeural'") {
		t.Errorf("ssml without default voice de: %q", ssml)
	}
	if !strings.Contains(ssml, "Hello &amp; done") {
		t.Errorf("ssml without escaped text: %q", ssml)
	}
}

// TestEdgeLoopbackConnectionAbort: the server closes after the requests
// without audio → clear "protocol changed?" error.
func TestEdgeLoopbackConnectionAbort(t *testing.T) {
	ts, _ := fakeEdgeServer(t, nil, true)
	p := NewEdge(ChainConfig{Getenv: envMap(map[string]string{
		EnvEdgeEndpoint: "ws://" + ts.Listener.Addr().String(),
	})})
	_, err := p.Synthesize(context.Background(), "Hello", "de")
	if err == nil || !strings.Contains(err.Error(), "protocol") {
		t.Errorf("error = %v, want protocol-change hint", err)
	}
}

// TestEdgeIntegrationReal is the integration marker against the REAL
// Microsoft endpoint, skippable via testing.Short (free, but real network;
// standard test runs and -short skip it; one-off manual acceptance with
// TALK_TEST_REAL_EDGE=1).
func TestEdgeIntegrationReal(t *testing.T) {
	if testing.Short() {
		t.Skip("skipped (-short)")
	}
	if os.Getenv("TALK_TEST_REAL_EDGE") != "1" {
		t.Skip("skipped (real network); run with TALK_TEST_REAL_EDGE=1")
	}
	p := NewEdge(ChainConfig{})
	data, err := p.Synthesize(context.Background(), "Short speech test via Edge.", "de")
	if err != nil {
		t.Fatalf("Synthesize against real endpoint: %v", err)
	}
	if len(data) == 0 {
		t.Fatal("Synthesize against real endpoint: 0 bytes")
	}
	t.Logf("Edge integration: %d bytes of audio", len(data))
}