From ac04b7afb4baa66f5bd745c81391dd16b2464410 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sat, 30 Jan 2021 18:55:50 +0100
Subject: [PATCH] rewrite TestSource

---
 main_test.go | 260 ++++++++++++++++++++++-----------------------------
 1 file changed, 112 insertions(+), 148 deletions(-)

diff --git a/main_test.go b/main_test.go
index f479297e..2fdb13af 100644
--- a/main_test.go
+++ b/main_test.go
@@ -685,168 +685,132 @@ func TestAuthIpFail(t *testing.T) {
 	require.NotEqual(t, 0, cnt1.wait())
 }
 
-func TestSourceRtsp(t *testing.T) {
-	for _, proto := range []string{
-		"udp",
-		"tcp",
+func TestSource(t *testing.T) {
+	for _, source := range []string{
+		"rtsp_udp",
+		"rtsp_tcp",
+		"rtsps",
+		"rtmp",
 	} {
-		t.Run(proto, func(t *testing.T) {
-			p1, ok := testProgram("paths:\n" +
-				"  all:\n" +
-				"    readUser: testuser\n" +
-				"    readPass: testpass\n")
-			require.Equal(t, true, ok)
-			defer p1.close()
+		t.Run(source, func(t *testing.T) {
+			switch source {
+			case "rtsp_udp", "rtsp_tcp":
+				p1, ok := testProgram("rtspPort: 8555\n" +
+					"rtpPort: 8100\n" +
+					"rtcpPort: 8101\n" +
+					"paths:\n" +
+					"  all:\n" +
+					"    readUser: testuser\n" +
+					"    readPass: testpass\n")
+				require.Equal(t, true, ok)
+				defer p1.close()
+
+				cnt1, err := newContainer("ffmpeg", "source", []string{
+					"-re",
+					"-stream_loop", "-1",
+					"-i", "emptyvideo.ts",
+					"-c", "copy",
+					"-f", "rtsp",
+					"-rtsp_transport", "udp",
+					"rtsp://" + ownDockerIP + ":8555/teststream",
+				})
+				require.NoError(t, err)
+				defer cnt1.close()
+
+				p2, ok := testProgram("paths:\n" +
+					"  proxied:\n" +
+					"    source: rtsp://testuser:testpass@localhost:8555/teststream\n" +
+					"    sourceProtocol: " + source[len("rtsp_"):] + "\n" +
+					"    sourceOnDemand: yes\n")
+				require.Equal(t, true, ok)
+				defer p2.close()
+
+			case "rtsps":
+				serverCertFpath, err := writeTempFile(serverCert)
+				require.NoError(t, err)
+				defer os.Remove(serverCertFpath)
+
+				serverKeyFpath, err := writeTempFile(serverKey)
+				require.NoError(t, err)
+				defer os.Remove(serverKeyFpath)
+
+				p, ok := testProgram("rtspPort: 8555\n" +
+					"rtpPort: 8100\n" +
+					"rtcpPort: 8101\n" +
+					"readTimeout: 20s\n" +
+					"protocols: [tcp]\n" +
+					"encryption: yes\n" +
+					"serverCert: " + serverCertFpath + "\n" +
+					"serverKey: " + serverKeyFpath + "\n" +
+					"paths:\n" +
+					"  all:\n" +
+					"    readUser: testuser\n" +
+					"    readPass: testpass\n")
+				require.Equal(t, true, ok)
+				defer p.close()
+
+				cnt1, err := newContainer("ffmpeg", "source", []string{
+					"-re",
+					"-stream_loop", "-1",
+					"-i", "emptyvideo.ts",
+					"-c", "copy",
+					"-f", "rtsp",
+					"rtsps://" + ownDockerIP + ":8555/teststream",
+				})
+				require.NoError(t, err)
+				defer cnt1.close()
+
+				time.Sleep(1 * time.Second)
+
+				p2, ok := testProgram("paths:\n" +
+					"  proxied:\n" +
+					"    source: rtsps://testuser:testpass@localhost:8555/teststream\n" +
+					"    sourceOnDemand: yes\n")
+				require.Equal(t, true, ok)
+				defer p2.close()
+
+			case "rtmp":
+				cnt1, err := newContainer("nginx-rtmp", "rtmpserver", []string{})
+				require.NoError(t, err)
+				defer cnt1.close()
+
+				cnt2, err := newContainer("ffmpeg", "source", []string{
+					"-re",
+					"-stream_loop", "-1",
+					"-i", "emptyvideo.ts",
+					"-c", "copy",
+					"-f", "flv",
+					"rtmp://" + cnt1.ip() + "/stream/test",
+				})
+				require.NoError(t, err)
+				defer cnt2.close()
+
+				time.Sleep(1 * time.Second)
+
+				p, ok := testProgram("paths:\n" +
+					"  proxied:\n" +
+					"    source: rtmp://" + cnt1.ip() + "/stream/test\n" +
+					"    sourceOnDemand: yes\n")
+				require.Equal(t, true, ok)
+				defer p.close()
+			}
 
 			time.Sleep(1 * time.Second)
 
-			cnt1, err := newContainer("ffmpeg", "source", []string{
-				"-re",
-				"-stream_loop", "-1",
-				"-i", "emptyvideo.ts",
-				"-c", "copy",
-				"-f", "rtsp",
+			cnt3, err := newContainer("ffmpeg", "dest", []string{
 				"-rtsp_transport", "udp",
-				"rtsp://" + ownDockerIP + ":8554/teststream",
-			})
-			require.NoError(t, err)
-			defer cnt1.close()
-
-			time.Sleep(1 * time.Second)
-
-			p2, ok := testProgram("rtspPort: 8555\n" +
-				"rtpPort: 8100\n" +
-				"rtcpPort: 8101\n" +
-				"\n" +
-				"paths:\n" +
-				"  proxied:\n" +
-				"    source: rtsp://testuser:testpass@localhost:8554/teststream\n" +
-				"    sourceProtocol: " + proto + "\n" +
-				"    sourceOnDemand: yes\n")
-			require.Equal(t, true, ok)
-			defer p2.close()
-
-			time.Sleep(1 * time.Second)
-
-			cnt2, err := newContainer("ffmpeg", "dest", []string{
-				"-rtsp_transport", "udp",
-				"-i", "rtsp://" + ownDockerIP + ":8555/proxied",
+				"-i", "rtsp://" + ownDockerIP + ":8554/proxied",
 				"-vframes", "1",
 				"-f", "image2",
 				"-y", "/dev/null",
 			})
 			require.NoError(t, err)
-			defer cnt2.close()
-
-			require.Equal(t, 0, cnt2.wait())
+			defer cnt3.close()
+			require.Equal(t, 0, cnt3.wait())
 		})
 	}
 }
 
-func TestSourceRtsps(t *testing.T) {
-	serverCertFpath, err := writeTempFile(serverCert)
-	require.NoError(t, err)
-	defer os.Remove(serverCertFpath)
-
-	serverKeyFpath, err := writeTempFile(serverKey)
-	require.NoError(t, err)
-	defer os.Remove(serverKeyFpath)
-
-	p, ok := testProgram("readTimeout: 20s\n" +
-		"protocols: [tcp]\n" +
-		"encryption: yes\n" +
-		"serverCert: " + serverCertFpath + "\n" +
-		"serverKey: " + serverKeyFpath + "\n" +
-		"paths:\n" +
-		"  all:\n" +
-		"    readUser: testuser\n" +
-		"    readPass: testpass\n")
-	require.Equal(t, true, ok)
-	defer p.close()
-
-	time.Sleep(1 * time.Second)
-
-	cnt1, err := newContainer("ffmpeg", "source", []string{
-		"-re",
-		"-stream_loop", "-1",
-		"-i", "emptyvideo.ts",
-		"-c", "copy",
-		"-f", "rtsp",
-		"rtsps://" + ownDockerIP + ":8555/teststream",
-	})
-	require.NoError(t, err)
-	defer cnt1.close()
-
-	time.Sleep(1 * time.Second)
-
-	p2, ok := testProgram("rtspPort: 8556\n" +
-		"rtpPort: 8100\n" +
-		"rtcpPort: 8101\n" +
-		"\n" +
-		"paths:\n" +
-		"  proxied:\n" +
-		"    source: rtsps://testuser:testpass@localhost:8555/teststream\n" +
-		"    sourceOnDemand: yes\n")
-	require.Equal(t, true, ok)
-	defer p2.close()
-
-	time.Sleep(1 * time.Second)
-
-	cnt2, err := newContainer("ffmpeg", "dest", []string{
-		"-rtsp_transport", "udp",
-		"-i", "rtsp://" + ownDockerIP + ":8556/proxied",
-		"-vframes", "1",
-		"-f", "image2",
-		"-y", "/dev/null",
-	})
-	require.NoError(t, err)
-	defer cnt2.close()
-
-	require.Equal(t, 0, cnt2.wait())
-}
-
-func TestSourceRtmp(t *testing.T) {
-	cnt1, err := newContainer("nginx-rtmp", "rtmpserver", []string{})
-	require.NoError(t, err)
-	defer cnt1.close()
-
-	time.Sleep(1 * time.Second)
-
-	cnt2, err := newContainer("ffmpeg", "source", []string{
-		"-re",
-		"-stream_loop", "-1",
-		"-i", "emptyvideo.ts",
-		"-c", "copy",
-		"-f", "flv",
-		"rtmp://" + cnt1.ip() + "/stream/test",
-	})
-	require.NoError(t, err)
-	defer cnt2.close()
-
-	time.Sleep(1 * time.Second)
-
-	p, ok := testProgram("paths:\n" +
-		"  proxied:\n" +
-		"    source: rtmp://" + cnt1.ip() + "/stream/test\n" +
-		"    sourceOnDemand: yes\n")
-	require.Equal(t, true, ok)
-	defer p.close()
-
-	time.Sleep(1 * time.Second)
-
-	cnt3, err := newContainer("ffmpeg", "dest", []string{
-		"-rtsp_transport", "udp",
-		"-i", "rtsp://" + ownDockerIP + ":8554/proxied",
-		"-vframes", "1",
-		"-f", "image2",
-		"-y", "/dev/null",
-	})
-	require.NoError(t, err)
-	defer cnt3.close()
-
-	require.Equal(t, 0, cnt3.wait())
-}
-
 func TestRedirect(t *testing.T) {
 	p1, ok := testProgram("paths:\n" +
 		"  path1:\n" +