speed up tests

main_rtspsource_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"crypto/tls"
 	"os"
 	"testing"
 	"time"
@@ -11,6 +12,57 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type testServer struct {
+	user          string
+	pass          string
+	authValidator *auth.Validator
+	done          chan struct{}
+}
+
+func (sh *testServer) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, []byte, error) {
+	if sh.authValidator == nil {
+		sh.authValidator = auth.NewValidator(sh.user, sh.pass, nil)
+	}
+
+	err := sh.authValidator.ValidateHeader(ctx.Req.Header["Authorization"],
+		ctx.Req.Method, ctx.Req.URL, nil)
+	if err != nil {
+		return &base.Response{
+			StatusCode: base.StatusUnauthorized,
+			Header: base.Header{
+				"WWW-Authenticate": sh.authValidator.GenerateHeader(),
+			},
+		}, nil, nil
+	}
+
+	track, _ := gortsplib.NewTrackH264(96, []byte{0x01, 0x02, 0x03, 0x04}, []byte{0x05, 0x06})
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, gortsplib.Tracks{track}.Write(), nil
+}
+
+func (sh *testServer) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *uint32, error) {
+	if sh.done != nil {
+		close(sh.done)
+	}
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, nil, nil
+}
+
+func (sh *testServer) OnPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
+	go func() {
+		time.Sleep(1 * time.Second)
+		ctx.Session.WriteFrame(0, gortsplib.StreamTypeRTP, []byte{0x01, 0x02, 0x03, 0x04})
+	}()
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, nil
+}
+
 func TestSourceRTSP(t *testing.T) {
 	for _, source := range []string{
 		"udp",
@@ -18,41 +70,14 @@ func TestSourceRTSP(t *testing.T) {
 		"tls",
 	} {
 		t.Run(source, func(t *testing.T) {
-			switch source {
-			case "udp", "tcp":
-				p1, ok := testProgram("rtmpDisable: yes\n" +
-					"hlsDisable: yes\n" +
-					"rtspAddress: :8555\n" +
-					"rtpAddress: :8100\n" +
-					"rtcpAddress: :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.mkv",
-					"-c", "copy",
-					"-f", "rtsp",
-					"-rtsp_transport", "udp",
-					"rtsp://" + ownDockerIP + ":8555/teststream",
-				})
-				require.NoError(t, err)
-				defer cnt1.close()
+			s := gortsplib.Server{
+				Handler: &testServer{user: "testuser", pass: "testpass"},
+			}
 
-				p2, ok := testProgram("rtmpDisable: yes\n" +
-					"hlsDisable: yes\n" +
-					"paths:\n" +
-					"  proxied:\n" +
-					"    source: rtsp://testuser:testpass@localhost:8555/teststream\n" +
-					"    sourceProtocol: " + source[len(""):] + "\n" +
-					"    sourceOnDemand: yes\n")
-				require.Equal(t, true, ok)
-				defer p2.close()
+			switch source {
+			case "udp":
+				s.UDPRTPAddress = "127.0.0.1:8002"
+				s.UDPRTCPAddress = "127.0.0.1:8003"
 
 			case "tls":
 				serverCertFpath, err := writeTempFile(serverCert)
@@ -63,102 +88,59 @@ func TestSourceRTSP(t *testing.T) {
 				require.NoError(t, err)
 				defer os.Remove(serverKeyFpath)
 
-				p, ok := testProgram("rtmpDisable: yes\n" +
-					"hlsDisable: yes\n" +
-					"rtspAddress: :8555\n" +
-					"rtpAddress: :8100\n" +
-					"rtcpAddress: :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.mkv",
-					"-c", "copy",
-					"-f", "rtsp",
-					"rtsps://" + ownDockerIP + ":8555/teststream",
-				})
+				cert, err := tls.LoadX509KeyPair(serverCertFpath, serverKeyFpath)
 				require.NoError(t, err)
-				defer cnt1.close()
 
-				time.Sleep(1 * time.Second)
+				s.TLSConfig = &tls.Config{Certificates: []tls.Certificate{cert}}
+			}
+
+			err := s.Start("127.0.0.1:8555")
+			require.NoError(t, err)
+			defer s.Close()
 
-				p2, ok := testProgram("rtmpDisable: yes\n" +
-					"hlsDisable: yes\n" +
-					"paths:\n" +
+			if source == "udp" || source == "tcp" {
+				p, ok := testProgram("paths:\n" +
+					"  proxied:\n" +
+					"    source: rtsp://testuser:testpass@localhost:8555/teststream\n" +
+					"    sourceProtocol: " + source + "\n" +
+					"    sourceOnDemand: yes\n")
+				require.Equal(t, true, ok)
+				defer p.close()
+			} else {
+				p, ok := testProgram("paths:\n" +
 					"  proxied:\n" +
 					"    source: rtsps://testuser:testpass@localhost:8555/teststream\n" +
 					"    sourceFingerprint: 33949E05FFFB5FF3E8AA16F8213A6251B4D9363804BA53233C4DA9A46D6F2739\n" +
 					"    sourceOnDemand: yes\n")
 				require.Equal(t, true, ok)
-				defer p2.close()
+				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",
-			})
+			conn, err := gortsplib.DialRead("rtsp://127.0.0.1:8554/proxied")
 			require.NoError(t, err)
-			defer cnt3.close()
-			require.Equal(t, 0, cnt3.wait())
-		})
-	}
-}
 
-type testServerNoPassword struct {
-	authValidator *auth.Validator
-	done          chan struct{}
-}
-
-func (sh *testServerNoPassword) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, []byte, error) {
-	if sh.authValidator == nil {
-		sh.authValidator = auth.NewValidator("testuser", "", nil)
-	}
+			readDone := make(chan struct{})
+			received := make(chan struct{})
+			go func() {
+				defer close(readDone)
+				conn.ReadFrames(func(trackID int, streamType gortsplib.StreamType, payload []byte) {
+					require.Equal(t, []byte{0x01, 0x02, 0x03, 0x04}, payload)
+					close(received)
+				})
+			}()
 
-	err := sh.authValidator.ValidateHeader(ctx.Req.Header["Authorization"],
-		ctx.Req.Method, ctx.Req.URL, nil)
-	if err != nil {
-		return &base.Response{
-			StatusCode: base.StatusUnauthorized,
-			Header: base.Header{
-				"WWW-Authenticate": sh.authValidator.GenerateHeader(),
-			},
-		}, nil, nil
+			<-received
+			conn.Close()
+			<-readDone
+		})
 	}
-
-	track, _ := gortsplib.NewTrackH264(96, []byte{0x01, 0x02, 0x03, 0x04}, []byte{0x05, 0x06})
-
-	return &base.Response{
-		StatusCode: base.StatusOK,
-	}, gortsplib.Tracks{track}.Write(), nil
-}
-
-// called after receiving a SETUP request.
-func (sh *testServerNoPassword) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *uint32, error) {
-	close(sh.done)
-	return &base.Response{
-		StatusCode: base.StatusOK,
-	}, nil, nil
 }
 
 func TestSourceRTSPNoPassword(t *testing.T) {
 	done := make(chan struct{})
-	s := gortsplib.Server{Handler: &testServerNoPassword{done: done}}
+	s := gortsplib.Server{Handler: &testServer{user: "testuser", done: done}}
 	err := s.Start("127.0.0.1:8555")
 	require.NoError(t, err)
 	defer s.Close()
@@ -173,5 +155,4 @@ func TestSourceRTSPNoPassword(t *testing.T) {
 	defer p.close()
 
 	<-done
-	// require.Equal(t, 0, cnt.wait())
 }