ffmpeg_test.go

package ffmpeg_go

import (
	"bytes"
	"fmt"
	"log"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/u2takey/go-utils/rand"
)

const (
	TestInputFile1  = "./examples/sample_data/in1.mp4"
	TestOutputFile1 = "./examples/sample_data/out1.mp4"
	TestOverlayFile = "./examples/sample_data/overlay.png"
)

// ffmpeg -y -rtsp_transport tcp -an -i rtsp://admin:password01!@192.168.188.21:554/Streaming/Channels/601
// -r 15 -g 30 -c:v copy -c:a copy -b:v 1000k -flags +global_header -strict experimental
// -f tee -map 0:v
// [f=segment:segment_time=\'00:01:00\':reset_timestamps=1:segment_atclocktime=1:strftime=1:segment_format=mpegts]
// static/records/8e7bf4ddff2a4f18adb36496334356fe/%Y-%m-%d-%H-%M-%S.ts|
// [f=flv]rtmp://192.168.111.140/live/8e7bf4ddff2a4f18adb36496334356fe

// ffmpeg -y -an -i rtmp://127.0.0.1/live/test -c:v copy -c:a copy -flags +global_header -strict experimental -f tee
// -map 0:v "[f=segment:reset_timestamps=1:segment_atclocktime=1:segment_time=60:strftime=1]./record/%Y%m%d%H%M.mp4|[f=flv]rtmp://127.0.0.1/live/mac"

// ffmpeg -use_wallclock_as_timestamps 1 - rtsp_transport tcp -i rtsp://[username]:[password]@[ip]:[port]
// -vcodec copy -acodec copy -flags +global_header -strict experimental
// -f segment -reset_timestamps 1 -segment_atclocktime 1 -segment_time 60 -strftime 1 /outdir/%Y%m%d%H%M.mp4

// ffmpeg -y -an -i rtmp://127.0.0.1/live/test -c:v copy -c:a copy  -flags +global_header -strict experimental
// -f tee -map 0:v "[f=segment -reset_timestamps=1 -segment_atclocktime=1 -segment_time=60 -strftime=1]./record/%Y%m%d%H%M.mp4|[f=flv]rtmp://127.0.0.1/live/mac"

// ffmpeg -y -an -i rtmp://127.0.0.1/live/test -c:v copy -c:a copy  -flags +global_header -strict experimental
// -f tee -map 0:v "[f=segment:reset_timestamps=1:segment_atclocktime=1:segment_time=60:strftime=1]./record/%Y%m%d%H%M.mp4|[f=flv]rtmp://127.0.0.1/live/mac"
// multi
func TestMultiOutput(t *testing.T) {
	Input("rtmp://127.0.0.1/live/test").
		Output("[f=segment:reset_timestamps=1:segment_atclocktime=1:segment_time=60:strftime=1]"+
			"/Users/zhangguoqing/github/ffmpeg-go/record/%Y%m%d%H%M.mp4|"+
			"[f=flv]rtmp://127.0.0.1/live/mac1",
			KwArgs{
				"flags":  "+global_header",
				"strict": "experimental",
			},
			KwArgs{
				"vcodec": "copy",
				"acodec": "copy",
			},
			KwArgs{
				"f":   "tee",
				"map": "0:v",
			}).GlobalArgs("-y", "-an").Run(nil)
	// log.Println(output.GetArgs())
	// output.String()
}

func TestFluentEquality(t *testing.T) {
	base1 := Input("dummy1.mp4")
	base2 := Input("dummy1.mp4")
	base3 := Input("dummy2.mp4")
	t1 := base1.Trim(KwArgs{"start_frame": 10, "end_frame": 20})
	t2 := base1.Trim(KwArgs{"start_frame": 10, "end_frame": 20})
	t3 := base1.Trim(KwArgs{"start_frame": 10, "end_frame": 30})
	t4 := base2.Trim(KwArgs{"start_frame": 10, "end_frame": 20})
	t5 := base3.Trim(KwArgs{"start_frame": 10, "end_frame": 20})

	assert.Equal(t, t1.Hash(), t2.Hash())
	assert.Equal(t, t1.Hash(), t4.Hash())
	assert.NotEqual(t, t1.Hash(), t3.Hash())
	assert.NotEqual(t, t1.Hash(), t5.Hash())
}

func TestFluentConcat(t *testing.T) {
	base1 := Input("dummy1.mp4", nil)
	trim1 := base1.Trim(KwArgs{"start_frame": 10, "end_frame": 20})
	trim2 := base1.Trim(KwArgs{"start_frame": 30, "end_frame": 40})
	trim3 := base1.Trim(KwArgs{"start_frame": 50, "end_frame": 60})
	concat1 := Concat([]*Stream{trim1, trim2, trim3})
	concat2 := Concat([]*Stream{trim1, trim2, trim3})
	concat3 := Concat([]*Stream{trim1, trim3, trim2})
	assert.Equal(t, concat1.Hash(), concat2.Hash())
	assert.NotEqual(t, concat1.Hash(), concat3.Hash())
}

func TestRepeatArgs(t *testing.T) {
	o := Input("dummy.mp4", nil).Output("dummy2.mp4",
		KwArgs{"streamid": []string{"0:0x101", "1:0x102"}})
	assert.Equal(t, o.GetArgs(), []string{"-i", "dummy.mp4", "-streamid", "0:0x101", "-streamid", "1:0x102", "dummy2.mp4"})
}

func TestGlobalArgs(t *testing.T) {
	o := Input("dummy.mp4", nil).Output("dummy2.mp4", nil).GlobalArgs("-progress", "someurl")

	assert.Equal(t, o.GetArgs(), []string{
		"-i",
		"dummy.mp4",
		"dummy2.mp4",
		"-progress",
		"someurl",
	})
}

func TestSimpleExample(t *testing.T) {
	err := Input(TestInputFile1, nil).
		Output(TestOutputFile1, nil).
		OverWriteOutput().
		Run(nil)
	assert.Nil(t, err)
}

func TestSimpleOverLayExample(t *testing.T) {
	s := Input(TestInputFile1, nil).Overlay(Input(TestOverlayFile), "").
		Output(TestOutputFile1).OverWriteOutput()
	err := s.Run(func() {
		log.Println("pid ", s.Context.Value("pid").(int))
	})

	// this code will panic, because s is not complete
	// s := Input(TestInputFile1, nil)
	// err := s.Overlay(Input(TestOverlayFile), "").
	//	Output(TestOutputFile1).OverWriteOutput().Run(func() {
	//	log.Println("pid ", s.Context.Value("pid").(int))
	// })
	assert.Nil(t, err)
}

func TestSimpleOutputArgs(t *testing.T) {
	cmd := Input(TestInputFile1).Output("imageFromVideo_%d.jpg", KwArgs{"vf": "fps=3", "qscale:v": 2})
	assert.Equal(t, []string{
		"-i", "./examples/sample_data/in1.mp4", "-qscale:v",
		"2", "-vf", "fps=3",
		"imageFromVideo_%d.jpg",
	}, cmd.GetArgs())
}

func TestAutomaticStreamSelection(t *testing.T) {
	// example from http://ffmpeg.org/ffmpeg-all.html
	input := []*Stream{Input("A.avi"), Input("B.mp4")}
	out1 := Output(input, "out1.mkv")
	out2 := Output(input, "out2.wav")
	out3 := Output(input, "out3.mov", KwArgs{"map": "1:a", "c:a": "copy"})
	cmd := MergeOutputs(out1, out2, out3)
	printArgs(cmd.GetArgs())
	printGraph(cmd)
}

func TestLabeledFiltergraph(t *testing.T) {
	// example from http://ffmpeg.org/ffmpeg-all.html
	in1, in2, in3 := Input("A.avi"), Input("B.mp4"), Input("C.mkv")
	in2Split := in2.Get("v").Hue(KwArgs{"s": 0}).Split()
	overlay := Filter([]*Stream{in1, in2}, "overlay", nil)
	aresample := Filter([]*Stream{in1, in2, in3}, "aresample", nil)
	out1 := Output([]*Stream{in2Split.Get("outv1"), overlay, aresample}, "out1.mp4", KwArgs{"an": ""})
	out2 := Output([]*Stream{in1, in2, in3}, "out2.mkv")
	out3 := in2Split.Get("outv2").Output("out3.mkv", KwArgs{"map": "1:a:0"})
	cmd := MergeOutputs(out1, out2, out3)
	printArgs(cmd.GetArgs())
	printGraph(cmd)
}

func ComplexFilterExample() *Stream {
	split := Input(TestInputFile1).VFlip().Split()
	split0, split1 := split.Get("0"), split.Get("1")
	overlayFile := Input(TestOverlayFile).Crop(10, 10, 158, 112)
	return Concat([]*Stream{
		split0.Trim(KwArgs{"start_frame": 10, "end_frame": 20}),
		split1.Trim(KwArgs{"start_frame": 30, "end_frame": 40}),
	}).
		Overlay(overlayFile.HFlip(), "").
		DrawBox(50, 50, 120, 120, "red", 5).
		Output(TestOutputFile1).
		OverWriteOutput()
}

func TestComplexFilterExample(t *testing.T) {
	assert.Equal(t, []string{
		"-i",
		TestInputFile1,
		"-i",
		TestOverlayFile,
		"-filter_complex",
		"[0]vflip[s0];" +
			"[s0]split=2[s1][s2];" +
			"[s1]trim=end_frame=20:start_frame=10[s3];" +
			"[s2]trim=end_frame=40:start_frame=30[s4];" +
			"[s3][s4]concat=n=2[s5];" +
			"[1]crop=158:112:10:10[s6];" +
			"[s6]hflip[s7];" +
			"[s5][s7]overlay=eof_action=repeat[s8];" +
			"[s8]drawbox=50:50:120:120:red:t=5[s9]",
		"-map",
		"[s9]",
		TestOutputFile1,
		"-y",
	}, ComplexFilterExample().GetArgs())
}

func TestCombinedOutput(t *testing.T) {
	i1 := Input(TestInputFile1)
	i2 := Input(TestOverlayFile)
	out := Output([]*Stream{i1, i2}, TestOutputFile1)
	assert.Equal(t, []string{
		"-i",
		TestInputFile1,
		"-i",
		TestOverlayFile,
		"-map",
		"0",
		"-map",
		"1",
		TestOutputFile1,
	}, out.GetArgs())
}

func TestFilterWithSelector(t *testing.T) {
	i := Input(TestInputFile1)

	v1 := i.Video().HFlip()
	a1 := i.Audio().Filter("aecho", Args{"0.8", "0.9", "1000", "0.3"})

	out := Output([]*Stream{a1, v1}, TestOutputFile1)
	assert.Equal(t, []string{
		"-i",
		TestInputFile1,
		"-filter_complex",
		"[0:a]aecho=0.8:0.9:1000:0.3[s0];[0:v]hflip[s1]",
		"-map",
		"[s0]",
		"-map",
		"[s1]",
		TestOutputFile1}, out.GetArgs())
}

func ComplexFilterAsplitExample() *Stream {
	split := Input(TestInputFile1).VFlip().ASplit()
	split0 := split.Get("0")
	split1 := split.Get("1")

	return Concat([]*Stream{
		split0.Filter("atrim", nil, KwArgs{"start": 10, "end": 20}),
		split1.Filter("atrim", nil, KwArgs{"start": 30, "end": 40}),
	}).Output(TestOutputFile1).OverWriteOutput()
}

func TestFilterConcatVideoOnly(t *testing.T) {
	in1 := Input("in1.mp4")
	in2 := Input("in2.mp4")
	args := Concat([]*Stream{in1, in2}).Output("out.mp4").GetArgs()
	assert.Equal(t, []string{
		"-i",
		"in1.mp4",
		"-i",
		"in2.mp4",
		"-filter_complex",
		"[0][1]concat=n=2[s0]",
		"-map",
		"[s0]",
		"out.mp4",
	}, args)
}

func TestConcatOutput(t *testing.T) {
	in1 := Input("in1.mp4")
	in2 := Input("in2.mp4")
	err := Concat([]*Stream{in1, in2}).Output("out.mp4").Run(nil)
	assert.Equal(t, nil, err)
}

func TestFilterConcatAudioOnly(t *testing.T) {
	in1 := Input("in1.mp4")
	in2 := Input("in2.mp4")
	args := Concat([]*Stream{in1, in2}, KwArgs{"v": 0, "a": 1}).Output("out.mp4").GetArgs()
	assert.Equal(t, []string{
		"-i",
		"in1.mp4",
		"-i",
		"in2.mp4",
		"-filter_complex",
		"[0][1]concat=a=1:n=2:v=0[s0]",
		"-map",
		"[s0]",
		"out.mp4",
	}, args)
}

func TestFilterConcatAudioVideo(t *testing.T) {
	in1 := Input("in1.mp4")
	in2 := Input("in2.mp4")
	joined := Concat([]*Stream{in1.Video(), in1.Audio(), in2.HFlip(), in2.Get("a")}, KwArgs{"v": 1, "a": 1}).Node
	args := Output([]*Stream{joined.Get("0"), joined.Get("1")}, "out.mp4").GetArgs()
	assert.Equal(t, []string{
		"-i",
		"in1.mp4",
		"-i",
		"in2.mp4",
		"-filter_complex",
		"[1]hflip[s0];[0:v][0:a][s0][1:a]concat=a=1:n=2:v=1[s1][s2]",
		"-map",
		"[s1]",
		"-map",
		"[s2]",
		"out.mp4",
	}, args)
}

func TestFilterASplit(t *testing.T) {
	out := ComplexFilterAsplitExample()
	args := out.GetArgs()
	assert.Equal(t, []string{
		"-i",
		TestInputFile1,
		"-filter_complex",
		"[0]vflip[s0];[s0]asplit=2[s1][s2];[s1]atrim=end=20:start=10[s3];[s2]atrim=end=40:start=30[s4];[s3][s4]concat=n=2[s5]",
		"-map",
		"[s5]",
		TestOutputFile1,
		"-y",
	}, args)
}

func TestOutputBitrate(t *testing.T) {
	args := Input("in").Output("out", KwArgs{"video_bitrate": 1000, "audio_bitrate": 200}).GetArgs()
	assert.Equal(t, []string{"-i", "in", "-b:v", "1000", "-b:a", "200", "out"}, args)
}

func TestOutputVideoSize(t *testing.T) {
	args := Input("in").Output("out", KwArgs{"video_size": "320x240"}).GetArgs()
	assert.Equal(t, []string{"-i", "in", "-video_size", "320x240", "out"}, args)
}

func TestCompile(t *testing.T) {
	out := Input("dummy.mp4").Output("dummy2.mp4")
	assert.Equal(t, out.Compile().Args, []string{"ffmpeg", "-i", "dummy.mp4", "dummy2.mp4"})
}

//	func TestCompileWithOptions(t *testing.T) {
//		out := Input("dummy.mp4").Output("dummy2.mp4")
//		cmd := out.Compile(SeparateProcessGroup())
//		assert.Equal(t, cmd.SysProcAttr.Pgid, 0)
//		assert.True(t, cmd.SysProcAttr.Setpgid)
//	}
func TestPipe(t *testing.T) {
	width, height := 32, 32
	frameSize := width * height * 3
	frameCount, startFrame := 10, 2
	_, _ = frameCount, frameSize

	out := Input(
		"pipe:0",
		KwArgs{
			"format":       "rawvideo",
			"pixel_format": "rgb24",
			"video_size":   fmt.Sprintf("%dx%d", width, height),
			"framerate":    10}).
		Trim(KwArgs{"start_frame": startFrame}).
		Output("pipe:1", KwArgs{"format": "rawvideo"})

	args := out.GetArgs()
	assert.Equal(t, args, []string{
		"-f",
		"rawvideo",
		"-video_size",
		fmt.Sprintf("%dx%d", width, height),
		"-framerate",
		"10",
		"-pixel_format",
		"rgb24",
		"-i",
		"pipe:0",
		"-filter_complex",
		"[0]trim=start_frame=2[s0]",
		"-map",
		"[s0]",
		"-f",
		"rawvideo",
		"pipe:1",
	})

	inBuf := bytes.NewBuffer(nil)
	for i := 0; i < frameSize*frameCount; i++ {
		inBuf.WriteByte(byte(rand.IntnRange(0, 255)))
	}
	outBuf := bytes.NewBuffer(nil)
	err := out.WithInput(inBuf).WithOutput(outBuf).Run(nil)
	assert.Nil(t, err)
	assert.Equal(t, outBuf.Len(), frameSize*(frameCount-startFrame))
}

func TestView(t *testing.T) {
	a, err := ComplexFilterExample().View(ViewTypeFlowChart)
	assert.Nil(t, err)

	b, err := ComplexFilterAsplitExample().View(ViewTypeStateDiagram)
	assert.Nil(t, err)

	t.Log(a)
	t.Log(b)
}

func TestInputOrder(t *testing.T) {
	const streamNum = 12
	streams := make([]*Stream, 0, streamNum)
	expectedFilenames := make([]string, 0, streamNum)
	for i := 0; i < streamNum; i++ {
		filename := fmt.Sprintf("%02d.mp4", i) // These files don't exist, but it's fine in this test.
		expectedFilenames = append(expectedFilenames, filename)
		streams = append(streams, Input(filename))
	}
	stream := Concat(streams).Output("output.mp4")
	args := stream.GetArgs()
	actualFilenames := make([]string, 0, streamNum)
	// Command looks like `-i 00.mp4 -i 01.mp4...` so actual filenames are the first streamNum even args
	for i := 1; i <= streamNum*2; i += 2 {
		actualFilenames = append(actualFilenames, args[i])
	}
	assert.Equal(t, expectedFilenames, actualFilenames)
}

func printArgs(args []string) {
	for _, a := range args {
		fmt.Printf("%s ", a)
	}
	fmt.Println()
}

func printGraph(s *Stream) {
	fmt.Println()
	v, _ := s.View(ViewTypeFlowChart)
	fmt.Println(v)
}