Piik/internal/app/nativehost/session_test.go

package nativehost

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"os"
	"slices"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/TNTcraftHIM/Piik/internal/app/mediaedge"
	"github.com/TNTcraftHIM/Piik/internal/app/nativecapture"
)

func TestMain(tests *testing.M) {
	if os.Getenv("PIIK_NATIVEHOST_PIPE_FIXTURE") == "backend" {
		for _, key := range []string{"PIIK_CAPTURE_STARTING", "PIIK_CAPTURE_ACTIVE"} {
			writeCaptureFixtureFrame(nativecapture.Frame{Kind: nativecapture.FrameStatus, Data: []byte(os.Getenv(key))})
			if layer := os.Getenv("PIIK_CAPTURE_UNAVAILABLE"); key == "PIIK_CAPTURE_STARTING" && layer != "" {
				value, _ := strconv.Atoi(layer)
				writeCaptureFixtureFrame(nativecapture.Frame{Kind: nativecapture.FrameLayerUnavailable, Layer: value, Data: []byte("fixture")})
			}
		}
		_, _ = io.Copy(io.Discard, os.Stdin)
		os.Exit(0)
	}
	if os.Getenv("PIIK_NATIVEHOST_PIPE_FIXTURE") == "audio-recovery" {
		runAudioRecoveryCapture()
		os.Exit(0)
	}
	if os.Getenv("PIIK_NATIVEHOST_PIPE_FIXTURE") == "echo" {
		_, _ = io.Copy(os.Stdout, os.Stdin)
		os.Exit(0)
	}
	if os.Getenv("PIIK_NATIVEHOST_PIPE_FIXTURE") == "1" {
		_, _ = io.Copy(io.Discard, os.Stdin)
		os.Exit(0)
	}
	os.Exit(tests.Run())
}

func TestCaptureCommitWaitsForReaderMetadataOrTermination(t *testing.T) {
	t.Setenv("PIIK_NATIVEHOST_PIPE_FIXTURE", "1")
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	for _, mode := range []string{"installed", "reader-failed", "closed"} {
		t.Run(mode, func(t *testing.T) {
			check := func(err error) {
				t.Helper()
				if err != nil {
					t.Fatal(err)
				}
			}
			ctx, cancel := context.WithCancel(t.Context())
			defer cancel()
			engine, err := mediaedge.NewEngine(mediaedge.EngineOptions{BindAddress: "127.0.0.1:0", IncludeLoopback: true})
			check(err)
			defer engine.Close()
			options := nativecapture.VideoOptions{Target: nativecapture.CaptureTarget{Kind: "display", SourceID: "1", Title: "Fixture"}, Codec: "vp8",
				Profile: nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 15, Bitrate: 2_000_000, Preference: "balanced"}, OutputGroups: 1}
			previous, err := nativecapture.StartVideo(ctx, executable, options)
			check(err)
			defer previous.Close()
			if len(previous.Outputs()) != 3 {
				t.Fatal("capture did not retain its independent output group")
			}
			source, err := engine.NewSource("vp8", 1, 2, nil)
			check(err)
			defer source.Close()
			check(configureCaptureOutputs(source, previous.Outputs()))
			publication, err := engine.NewPublication(source, mediaedge.EdgeOptions{ConnectionID: "metadata-fixture"})
			check(err)
			session := &Session{ctx: ctx, cancel: cancel, engine: engine, source: source, stream: previous,
				edgeCapacity: 1, ready: make(chan error, 1), done: make(chan error, 1), events: func(context.Context, Event) {}}
			options.Profile.Width, options.Profile.Height = 854, 480
			replacement, err := nativecapture.StartVideo(ctx, executable, options)
			check(err)
			defer replacement.Close()
			committed := make(chan error, 1)
			adapter, encoder := uint32(1), uint32(2)
			go func() {
				session.updateMu.Lock()
				err := session.commitCapture(options, QualityProfile{Video: options.Profile, AudioBitrate: 64_000}, replacement,
					CaptureState{AdapterIndex: &adapter, EncoderIndex: &encoder, unavailableLayers: []int{0}}, nil, false)
				session.updateMu.Unlock()
				committed <- err
			}()
			select {
			case <-previous.Done():
			case <-time.After(5 * time.Second):
				t.Fatal("old fixture did not exit")
			}
			select {
			case err := <-committed:
				t.Fatalf("process exit acknowledged metadata before reader installation: %v", err)
			default:
			}
			if len(publication.Media().Layers) != 2 || publication.Media().Layers[1].Width != 1280 {
				t.Fatal("held reader unexpectedly changed source metadata")
			}
			if mode == "installed" {
				// Hold the reader at its handoff, then run the same installation it uses.
				check(session.installCapture(replacement))
				select {
				case err = <-committed:
					check(err)
				case <-time.After(time.Second):
					t.Fatal("installed metadata did not acknowledge update")
				}
				media := publication.Media()
				if session.videoOptions.AdapterIndex != adapter || session.videoOptions.EncoderIndex != encoder {
					t.Fatal("capture replacement discarded its selected hardware backend")
				}
				if media.Layers[1].Width != 854 || media.Layers[1].Height != 480 || media.Layers[1].Bitrate != 2_000_000 {
					t.Fatalf("acknowledged stale metadata: %+v", media)
				}
				plan, err := source.BeginFrame(time.Second)
				check(err)
				if plan.Bitrates[0] != 0 || plan.Bitrates[1] == 0 || session.captureUnavailable != nil {
					t.Fatalf("replacement lost or retained pending failures: %+v", plan)
				}
				// The next healthy generation must not inherit the failed layer.
				healthy, err := nativecapture.StartVideo(ctx, executable, options)
				check(err)
				defer healthy.Close()
				go func() {
					committed <- session.commitCapture(options, session.profile, healthy, CaptureState{}, nil, false)
				}()
				select {
				case <-replacement.Done():
				case <-time.After(5 * time.Second):
					t.Fatal("replacement fixture did not exit")
				}
				check(session.installCapture(healthy))
				select {
				case err = <-committed:
					check(err)
				case <-time.After(time.Second):
					t.Fatal("healthy replacement remained pending")
				}
				plan, err = source.BeginFrame(time.Second)
				check(err)
				if plan.Bitrates[0] == 0 || plan.Bitrates[1] == 0 {
					t.Fatalf("healthy generation inherited a failed layer: %+v", plan)
				}
				return
			}
			go session.run()
			closed := make(chan error, 1)
			if mode == "closed" {
				go func() { closed <- session.Close() }()
			} else {
				check(replacement.Close())
			}
			select {
			case err = <-committed:
				if err == nil {
					t.Fatal("terminated reader acknowledged an uninstalled stream")
				}
			case <-time.After(5 * time.Second):
				t.Fatal("capture update remained latched after termination")
			}
			if mode == "closed" {
				select {
				case err = <-closed:
					check(err)
				case <-time.After(5 * time.Second):
					t.Fatal("Close deadlocked behind the capture update")
				}
			} else {
				select {
				case err = <-session.Done():
					if err == nil {
						t.Fatal("unexpected reader death did not fail the share")
					}
				case <-time.After(5 * time.Second):
					t.Fatal("reader death did not end the share")
				}
			}
		})
	}
}

func TestCaptureFailureBelongsToSelectedStream(t *testing.T) {
	for _, replace := range []bool{false, true} {
		t.Run(strconv.FormatBool(replace), func(t *testing.T) {
			check := func(err error) {
				t.Helper()
				if err != nil {
					t.Fatal(err)
				}
			}
			t.Setenv("PIIK_NATIVEHOST_PIPE_FIXTURE", "backend")
			t.Setenv("PIIK_CAPTURE_STARTING", string(captureStatePayload(t,
				`{"state":"starting","codec":"vp8","adapterIndex":0,"adapterName":"Fixture","adapterIdentity":"0:0","encoderName":"VP8","encoderIdentity":"vp8"}`)))
			t.Setenv("PIIK_CAPTURE_ACTIVE", string(captureStatePayload(t, `{"state":"active","codec":"vp8","width":1280,"height":720,"fps":30}`)))
			t.Setenv("PIIK_CAPTURE_UNAVAILABLE", "0")
			executable, err := os.Executable()
			check(err)
			ctx, cancel := context.WithCancel(t.Context())
			defer cancel()
			options := nativecapture.VideoOptions{Target: nativecapture.CaptureTarget{Kind: "display", SourceID: "1", Title: "Fixture"}, Codec: "vp8",
				Profile: nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 30, Bitrate: 3_000_000, Preference: "balanced"}, OutputGroups: 4}
			previous, err := nativecapture.StartVideo(ctx, executable, options)
			check(err)
			defer previous.Close()
			t.Setenv("PIIK_CAPTURE_UNAVAILABLE", "")
			replacement, err := nativecapture.StartVideo(ctx, executable, options)
			check(err)
			defer replacement.Close()
			engine, err := mediaedge.NewEngine(mediaedge.EngineOptions{BindAddress: "127.0.0.1:0", IncludeLoopback: true})
			check(err)
			defer engine.Close()
			source, err := engine.NewSource("vp8", 4, 2, nil)
			check(err)
			defer source.Close()
			check(configureCaptureOutputs(source, previous.Outputs()))
			publication, err := engine.NewPublication(source, mediaedge.EdgeOptions{ConnectionID: "capture-failure-owner"})
			check(err)
			active := make(chan struct{}, 1)
			session := &Session{ctx: ctx, cancel: cancel, engine: engine, source: source, stream: previous,
				videoOptions: options, ready: make(chan error, 1), events: func(_ context.Context, event Event) {
					if event.State == "active" {
						select {
						case active <- struct{}{}:
						default:
						}
					}
				}}
			// Delay the actual reader after its first identity check, without a
			// production hook. Log I/O and ordinary scheduling can delay it here.
			entered, release := make(chan struct{}), make(chan struct{})
			unblock := sync.OnceFunc(func() { close(release) })
			previousLog := slog.Default()
			defer slog.SetDefault(previousLog)
			slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{
				Level: slog.LevelDebug, ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
					if attr.Key == "event" && attr.Value.String() == "capture-output-unavailable" {
						close(entered)
						<-release
					}
					return attr
				},
			})))
			readerDone := make(chan error, 1)
			go func() { readerDone <- session.runVideo() }()
			defer func() {
				unblock()
				cancel()
				_ = previous.Close()
				_ = replacement.Close()
				select {
				case <-readerDone:
				case <-time.After(5 * time.Second):
					t.Error("capture reader did not retire")
				}
			}()
			select {
			case <-entered:
			case err = <-readerDone:
				readerDone <- err
				t.Fatalf("reader stopped before output failure: %v", err)
			case <-time.After(5 * time.Second):
				t.Fatal("old output failure did not reach the reader")
			}
			committed := make(chan error, 1)
			if replace {
				go func() {
					committed <- session.commitCapture(options, QualityProfile{Video: options.Profile}, replacement, CaptureState{}, nil, false)
				}()
				deadline := time.Now().Add(5 * time.Second)
				for session.currentStream() != replacement {
					if time.Now().After(deadline) {
						t.Fatal("replacement was not selected")
					}
					time.Sleep(time.Millisecond)
				}
			}
			unblock()
			select {
			case <-active:
			case <-time.After(5 * time.Second):
				t.Fatal("capture reader did not finish the failure or handoff")
			}
			if replace {
				select {
				case err = <-committed:
					check(err)
				case <-time.After(5 * time.Second):
					t.Fatal("replacement installation was not acknowledged")
				}
			}
			if err = publication.SetActiveCount(1); (err == nil) != replace {
				t.Fatalf("publication availability after replacement=%v: %v", replace, err)
			}
		})
	}
}

func TestCaptureProfileRetainsSelectedBackendFromStarting(t *testing.T) {
	t.Setenv("PIIK_NATIVEHOST_PIPE_FIXTURE", "backend")
	t.Setenv("PIIK_CAPTURE_STARTING", string(captureStatePayload(t,
		`{"state":"starting","hardwareOnly":true,"codec":"h264","adapterIndex":1,"adapterName":"GPU","adapterIdentity":"0:1","encoderIndex":2,"encoderName":"H264","encoderIdentity":"encoder"}`)))
	t.Setenv("PIIK_CAPTURE_ACTIVE", string(captureStatePayload(t,
		`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"42c01f","width":1280,"height":720,"fps":30}`)))
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	profile := nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 30, Bitrate: 3_000_000, Preference: "balanced"}
	for _, failed := range []string{"", "0", "1", "5", "6"} {
		t.Run("unavailable-"+failed, func(t *testing.T) {
			t.Setenv("PIIK_CAPTURE_UNAVAILABLE", failed)
			stream, err := nativecapture.StartVideo(t.Context(), executable, nativecapture.VideoOptions{
				Target: nativecapture.CaptureTarget{Kind: "display", SourceID: "1", Title: "Fixture"}, Codec: "h264", Profile: profile, OutputGroups: 4,
			})
			if err != nil {
				t.Fatal(err)
			}
			defer stream.Close()
			state, err := waitForCaptureProfile(t.Context(), stream, profile, "h264", false, false)
			if failed == "1" || failed == "6" {
				if err == nil {
					t.Fatal("original or invalid output failure was accepted")
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if state.AdapterIndex == nil || *state.AdapterIndex != 1 || state.EncoderIndex == nil || *state.EncoderIndex != 2 {
				t.Fatalf("active profile lost the actual hardware selection: %+v", state)
			}
			var unavailable []int
			if failed != "" {
				layer, _ := strconv.Atoi(failed)
				unavailable = []int{layer}
			}
			if !slices.Equal(state.unavailableLayers, unavailable) {
				t.Fatalf("preparation lost optional output failures: %+v", state.unavailableLayers)
			}
		})
	}
}

func TestOutputPlanPreservesOriginalStartupAndIndependentActivation(t *testing.T) {
	t.Setenv("PIIK_NATIVEHOST_PIPE_FIXTURE", "echo")
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
	defer cancel()
	stream, err := nativecapture.StartVideo(ctx, executable, nativecapture.VideoOptions{
		Target: nativecapture.CaptureTarget{Kind: "display", SourceID: "1", Title: "Fixture"}, Codec: "vp8",
		Profile: nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 15,
			Bitrate: 2_000_000, Preference: "balanced"}, OutputGroups: 2,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer stream.Close()
	plan := mediaedge.OutputPlan{Active: []bool{false, true, true, false}, Bitrates: []uint32{0, 0, 250_000, 0}}
	if err = applyOutputPlan(stream, plan, false); err != nil {
		t.Fatal(err)
	}
	plan.RecoveryLayers = []int{2}
	if err = applyOutputPlan(stream, plan, true); err != nil {
		t.Fatal(err)
	}
	clear(plan.Active)
	plan.RecoveryLayers = nil
	if err = applyOutputPlan(stream, plan, true); err != nil {
		t.Fatal(err)
	}
	for _, command := range []string{"B 2 250000", "K 2", "A 0 0", "A 2 1", "A 1 0", "A 2 0"} {
		frame, readErr := stream.Read()
		if readErr != nil || frame.Kind != nativecapture.FrameControl || string(frame.Data) != command {
			t.Fatalf("expected control %q, received %+v, %v", command, frame, readErr)
		}
	}
}

func TestCaptureStateKeepsStartingAndActiveContractsDistinct(t *testing.T) {
	starting, err := decodeCaptureState(captureStatePayload(t,
		`{"state":"starting","hardwareOnly":true,"codec":"h264","adapterIndex":0,"adapterName":"GPU","adapterIdentity":"0:1","encoderIndex":0,"encoderName":"H264","encoderIdentity":"encoder"}`,
	))
	if err != nil || starting.State != "starting" || starting.AdapterIndex == nil {
		t.Fatalf("starting = %+v, %v", starting, err)
	}
	active, err := decodeCaptureState(captureStatePayload(t,
		`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"42c01f","width":1280,"height":720,"fps":30}`,
	))
	if err != nil || active.State != "active" || active.ProfileLevelID != "42c01f" {
		t.Fatalf("active = %+v, %v", active, err)
	}
	for _, payload := range []string{
		`{"state":"starting","hardwareOnly":false,"codec":"vp8","adapterIndex":0,"adapterName":"GPU","adapterIdentity":"0:1","encoderName":"libvpx VP8","encoderIdentity":"libvpx/v1.17.0"}`,
		`{"state":"active","hardwareOnly":false,"codec":"vp8","width":1280,"height":720,"fps":30}`,
		`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"42c01e","width":854,"height":480,"fps":15}`,
		`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"42c033","width":2560,"height":1440,"fps":60,"restoreToken":"portal-token"}`,
	} {
		if _, err = decodeCaptureState(captureStatePayload(t, payload)); err != nil {
			t.Fatalf("valid profile state was rejected: %s: %v", payload, err)
		}
	}
}

func TestCaptureStateRejectsUnknownOrUnattributedState(t *testing.T) {
	for _, payload := range []string{
		`{"state":"active","hardwareOnly":true,"codec":"vp8","width":1280,"height":720,"fps":30}`,
		`{"state":"active","hardwareOnly":false,"codec":"vp8","profileLevelId":"42c01f","width":1280,"height":720,"fps":30}`,
		`{"state":"active","hardwareOnly":false}`,
		`{"state":"ready","hardwareOnly":true}`,
		`{"state":"active","hardwareOnly":true,"codec":"h264","extra":true}`,
	} {
		if _, err := decodeCaptureState(captureStatePayload(t, payload)); err == nil {
			t.Fatalf("invalid state was accepted: %s", payload)
		}
	}
}

func TestCaptureStateValidatesH264ToolsRatherThanConstraintSpelling(t *testing.T) {
	for _, profile := range []string{"42c01f", "42e01f", "42E033", "4d801f", "58c02a"} {
		state, err := decodeCaptureState(captureStatePayload(t,
			`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"`+profile+`","width":1280,"height":720,"fps":30}`))
		if err != nil || state.ProfileLevelID != profile {
			t.Fatalf("equivalent Constrained Baseline %q was rejected or rewritten: %+v, %v", profile, state, err)
		}
	}
	for _, profile := range []string{"42001f", "4d001f", "64001f", "42e11f", "42e01", "42e01fg", "42e0gg", "42e034", "42e035", "42e03c"} {
		if _, err := decodeCaptureState(captureStatePayload(t,
			`{"state":"active","hardwareOnly":true,"codec":"h264","profileLevelId":"`+profile+`","width":1280,"height":720,"fps":30}`)); err == nil {
			t.Fatalf("incompatible profile/level %q was admitted", profile)
		}
	}
}

func TestCaptureStateUsesOriginalSlotAndBoundsPhysicalGroups(t *testing.T) {
	state, err := decodeCaptureState(captureStatePayload(t,
		`{"state":"active","hardwareOnly":false,"codec":"vp8","width":1280,"height":720,"fps":30}`))
	if err != nil || len(state.Outputs) != 6 || state.Outputs[5] != state.Outputs[0] {
		t.Fatalf("independent output groups = %+v, %v", state, err)
	}
	wrongOriginal := state
	wrongOriginal.Width = state.Outputs[5].Width
	tooMany := state
	tooMany.Outputs = append(tooMany.Outputs, tooMany.Outputs[0])
	for _, invalid := range []CaptureState{wrongOriginal, tooMany} {
		payload, err := json.Marshal(invalid)
		if err != nil {
			t.Fatal(err)
		}
		if _, err = decodeCaptureState(payload); err == nil {
			t.Fatal("invalid original slot or output count was accepted")
		}
	}
}

func captureStatePayload(t *testing.T, payload string) []byte {
	t.Helper()
	var fields map[string]json.RawMessage
	if err := json.Unmarshal([]byte(payload), &fields); err != nil {
		t.Fatal(err)
	}
	profile := nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 30, Bitrate: 3_000_000}
	for key, field := range map[string]*uint32{"width": &profile.Width, "height": &profile.Height, "fps": &profile.Framerate} {
		if value, ok := fields[key]; ok {
			if err := json.Unmarshal(value, field); err != nil {
				t.Fatal(err)
			}
		}
	}
	slots := nativecapture.ScreenShareOutputs(profile)
	for range 4 {
		slots = append(slots, slots[0])
	}
	outputs, err := json.Marshal(slots)
	if err != nil {
		t.Fatal(err)
	}
	fields["outputs"] = outputs
	result, err := json.Marshal(fields)
	if err != nil {
		t.Fatal(err)
	}
	return result
}

func TestCandidateForRetiredEdgeIsIgnored(t *testing.T) {
	session := &Session{}
	if err := session.AddCandidate("retired-edge", nil); err != nil {
		t.Fatalf("stale candidate stopped the session: %v", err)
	}
}

func TestCaptureProfileFailurePreservesCauseWithoutCredentials(t *testing.T) {
	previous := slog.Default()
	t.Cleanup(func() { slog.SetDefault(previous) })
	var output bytes.Buffer
	slog.SetDefault(slog.New(slog.NewJSONHandler(&output, &slog.HandlerOptions{Level: slog.LevelDebug})))
	failure := errors.New("profile failed: hresult=0x887a0005 token=private-capture-token")
	profile := nativecapture.VideoProfile{Width: 1280, Height: 720, Framerate: 30, Bitrate: 3_000_000}
	if captureProfileFailure(t.Context(), profile, "wait-timeout", failure) != failure {
		t.Fatal("diagnostic changed the failure result")
	}
	var record map[string]any
	if err := json.Unmarshal(output.Bytes(), &record); err != nil {
		t.Fatal(err)
	}
	if record["event"] != "capture-profile-rejected" || record["stage"] != "wait-timeout" ||
		record["width"] != float64(1280) || record["height"] != float64(720) ||
		record["fps"] != float64(30) || record["bitrate"] != float64(3_000_000) ||
		!strings.Contains(output.String(), "0x887a0005") || strings.Contains(output.String(), "private-capture-token") {
		t.Fatalf("profile diagnostic changed or leaked: %s", output.String())
	}
}