linearvoid/backend/playout/src/playlist_rewrite_test.go

package main

import (
	"io"
	"os"
	"path/filepath"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
)

const sampleManifest = `#EXTM3U
#EXT-X-VERSION:7
#EXT-X-TARGETDURATION:2
#EXT-X-MEDIA-SEQUENCE:0
#EXT-X-INDEPENDENT-SEGMENTS
#EXT-X-MAP:URI="init.mp4"
#EXTINF:2.000000,
seg-000000.m4s
#EXTINF:2.000000,
seg-000001.m4s
#EXT-X-ENDLIST
`

func TestRewriteManifestToRun(t *testing.T) {
	runID := "r-20260903T120000Z"
	out := string(rewriteManifestToRun(runID, []byte(sampleManifest)))

	wants := []string{
		`#EXT-X-MAP:URI="r-20260903T120000Z/init.mp4"`,
		"r-20260903T120000Z/seg-000000.m4s",
		"r-20260903T120000Z/seg-000001.m4s",
		"#EXTM3U",
		"#EXT-X-VERSION:7",
		"#EXT-X-TARGETDURATION:2",
		"#EXT-X-MEDIA-SEQUENCE:0",
		"#EXT-X-INDEPENDENT-SEGMENTS",
		"#EXT-X-ENDLIST",
		"#EXTINF:2.000000,",
	}
	for _, want := range wants {
		if !contains(out, want) {
			t.Errorf("output missing %q\noutput:\n%s", want, out)
		}
	}

	// Media URIs must be prefixed (no bare segments left).
	for _, bare := range []string{"\nseg-000000.m4s\n", "\nseg-000001.m4s\n", `URI="init.mp4"`} {
		if contains(out, bare) {
			t.Errorf("output still contains bare %q\noutput:\n%s", bare, out)
		}
	}
}

func TestRewriteManifestToRun_AlreadyPrefixedUnchanged(t *testing.T) {
	manifest := "#EXTM3U\n#EXT-X-MAP:URI=\"r-old/init.mp4\"\nseg-000000.m4s\nr-old/seg-000001.m4s\n#EXT-X-ENDLIST\n"
	out := string(rewriteManifestToRun("r-new", []byte(manifest)))

	if !contains(out, `#EXT-X-MAP:URI="r-old/init.mp4"`) {
		t.Errorf("EXT-X-MAP already-prefixed URI was rewritten: %s", out)
	}
	// A URI line that already contains a path is left alone.
	if !contains(out, "r-old/seg-000001.m4s") {
		t.Errorf("already-prefixed segment URI was changed:\n%s", out)
	}
	// Bare URI still gets the run prefix.
	if !contains(out, "r-new/seg-000000.m4s") {
		t.Errorf("bare segment not prefixed:\n%s", out)
	}
}

func TestRewriteManifestToRun_HandlesTSAndComments(t *testing.T) {
	manifest := "#EXTM3U\n# a comment\nvideo-1.ts\ninit.mp4\n#EXT-X-ENDLIST\n"
	out := string(rewriteManifestToRun("r-x", []byte(manifest)))

	for _, want := range []string{"r-x/video-1.ts", "r-x/init.mp4", "# a comment", "#EXTM3U", "#EXT-X-ENDLIST"} {
		if !contains(out, want) {
			t.Errorf("output missing %q\noutput:\n%s", want, out)
		}
	}
}

func TestRewriteManifestToRun_EmptyRunID(t *testing.T) {
	out := rewriteManifestToRun("", []byte(sampleManifest))
	if string(out) != sampleManifest {
		t.Errorf("empty runID should return input unchanged")
	}
}

func TestIsMediaURI(t *testing.T) {
	for _, ok := range []string{"seg-000001.m4s", "init.mp4", "a.ts", "x.m4a"} {
		if !isMediaURI(ok) {
			t.Errorf("isMediaURI(%q) = false, want true", ok)
		}
	}
	for _, notOK := range []string{"playlist.m3u8", "readme.txt", "seg-000001.m4s?x=1"} {
		// Note: query strings are stripped; '?x=1' on m4s should still be media.
		if notOK == "seg-000001.m4s?x=1" {
			continue
		}
		if isMediaURI(notOK) {
			t.Errorf("isMediaURI(%q) = true, want false", notOK)
		}
	}
	if !isMediaURI("seg-000001.m4s?token=abc") {
		t.Errorf("isMediaURI with query should strip query and be true")
	}
}

func TestPublishManifest_RewritesToRunScopedURIs(t *testing.T) {
	prevClient, prevBucket := r2Client, bucket
	prevRunID, prevMode := runID, playoutMode
	prevPut, prevCopy := putObject, copyObject
	defer func() {
		r2Client, bucket = prevClient, prevBucket
		runID, playoutMode = prevRunID, prevMode
		putObject, copyObject = prevPut, prevCopy
	}()

	tmpDir := t.TempDir()
	mp := filepath.Join(tmpDir, "playlist.m3u8")
	if err := os.WriteFile(mp, []byte(sampleManifest), 0o644); err != nil {
		t.Fatalf("write manifest: %v", err)
	}

	bucket = "linearvoid-media"
	runID = "r-20260903T120000Z"
	playoutMode = "live"
	r2Client = &s3.Client{}

	var gotKey, gotCopySource, gotCopyKey string
	var gotBody []byte
	putObject = func(_ *s3.Client, in *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
		gotKey = aws.ToString(in.Key)
		b, err := io.ReadAll(in.Body)
		if err != nil {
			return nil, err
		}
		gotBody = b
		return &s3.PutObjectOutput{}, nil
	}
	copyObject = func(_ *s3.Client, in *s3.CopyObjectInput) (*s3.CopyObjectOutput, error) {
		gotCopySource = aws.ToString(in.CopySource)
		gotCopyKey = aws.ToString(in.Key)
		return &s3.CopyObjectOutput{}, nil
	}

	if err := publishManifest(mp); err != nil {
		t.Fatalf("publishManifest error: %v", err)
	}

	if gotKey != "hls/playlist.m3u8.tmp" {
		t.Errorf("temp upload key = %q, want hls/playlist.m3u8.tmp", gotKey)
	}
	if gotCopySource != "linearvoid-media/hls/playlist.m3u8.tmp" {
		t.Errorf("CopySource = %q, want linearvoid-media/hls/playlist.m3u8.tmp", gotCopySource)
	}
	if gotCopyKey != "hls/playlist.m3u8" {
		t.Errorf("final key = %q, want hls/playlist.m3u8", gotCopyKey)
	}

	body := string(gotBody)
	for _, want := range []string{
		"r-20260903T120000Z/seg-000000.m4s",
		"r-20260903T120000Z/seg-000001.m4s",
		`#EXT-X-MAP:URI="r-20260903T120000Z/init.mp4"`,
	} {
		if !contains(body, want) {
			t.Errorf("published manifest missing %q:\n%s", want, body)
		}
	}
	for _, bare := range []string{"\nseg-000000.m4s\n", "\nseg-000001.m4s\n"} {
		if contains(body, bare) {
			t.Errorf("published manifest still contains bare %q:\n%s", bare, body)
		}
	}
}

func TestPublishManifest_SlateNotRewritten(t *testing.T) {
	prevClient, prevBucket := r2Client, bucket
	prevRunID, prevMode := runID, playoutMode
	prevPut, prevCopy := putObject, copyObject
	defer func() {
		r2Client, bucket = prevClient, prevBucket
		runID, playoutMode = prevRunID, prevMode
		putObject, copyObject = prevPut, prevCopy
	}()

	tmpDir := t.TempDir()
	mp := filepath.Join(tmpDir, "playlist.m3u8")
	if err := os.WriteFile(mp, []byte(sampleManifest), 0o644); err != nil {
		t.Fatalf("write manifest: %v", err)
	}

	bucket = "linearvoid-media"
	runID = "r-20260903T120000Z"
	playoutMode = "slate"
	r2Client = &s3.Client{}

	var gotBody []byte
	putObject = func(_ *s3.Client, in *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
		b, _ := io.ReadAll(in.Body)
		gotBody = b
		return &s3.PutObjectOutput{}, nil
	}
	copyObject = func(_ *s3.Client, _ *s3.CopyObjectInput) (*s3.CopyObjectOutput, error) {
		return &s3.CopyObjectOutput{}, nil
	}

	if err := publishManifest(mp); err != nil {
		t.Fatalf("publishManifest error: %v", err)
	}

	if string(gotBody) != sampleManifest {
		t.Errorf("slate manifest should be published unmodified; got:\n%s", string(gotBody))
	}
}

func contains(s, sub string) bool {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return true
		}
	}
	return false
}