From 026c8725a3fe54271ff2f03b76525379e4a33a33 Mon Sep 17 00:00:00 2001
From: David Zhao <david@davidzhao.com>
Date: Tue, 22 Dec 2020 20:59:39 -0800
Subject: [PATCH] finished room test

---
 magefile.go                 | 26 +++++++++---------
 pkg/rtc/mock_helper_test.go |  8 ++++++
 pkg/rtc/room.go             |  4 +++
 pkg/rtc/room_test.go        | 53 ++++++++++++++++++++++++++++++++++---
 4 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/magefile.go b/magefile.go
index 559aaef31..60100e4b9 100644
--- a/magefile.go
+++ b/magefile.go
@@ -125,7 +125,7 @@ func Build() error {
 }
 
 func Test() error {
-	mg.Deps(Proto, generateTest)
+	mg.Deps(Proto)
 	cmd := exec.Command("go", "test", "./...")
 	connectStd(cmd)
 	return cmd.Run()
@@ -139,6 +139,17 @@ func Clean() {
 	os.Remove(goChecksumFile)
 }
 
+// regenerate code
+func Generate() error {
+	mg.Deps(installDeps)
+
+	fmt.Println("generating...")
+
+	cmd := exec.Command("go", "generate", "./...")
+	connectStd(cmd)
+	return cmd.Run()
+}
+
 // code generation
 func generateCmd() error {
 	mg.Deps(installDeps)
@@ -153,19 +164,6 @@ func generateCmd() error {
 	return cmd.Run()
 }
 
-func generateTest() error {
-	mg.Deps(installDeps)
-	if !checksummer.IsChanged() {
-		return nil
-	}
-
-	fmt.Println("generating for tests...")
-
-	cmd := exec.Command("go", "generate", "./pkg/...")
-	connectStd(cmd)
-	return cmd.Run()
-}
-
 // implicitly install deps
 func installDeps() error {
 	return installTools(false)
diff --git a/pkg/rtc/mock_helper_test.go b/pkg/rtc/mock_helper_test.go
index 89a92fd6f..a95aa216c 100644
--- a/pkg/rtc/mock_helper_test.go
+++ b/pkg/rtc/mock_helper_test.go
@@ -14,3 +14,11 @@ func newMockParticipant(name string) *rtcfakes.FakeParticipant {
 
 	return p
 }
+
+func newMockTrack(kind livekit.TrackInfo_Type, name string) *rtcfakes.FakePublishedTrack {
+	t := &rtcfakes.FakePublishedTrack{}
+	t.IDReturns(utils.NewGuid(utils.TrackPrefix))
+	t.KindReturns(kind)
+	t.StreamIDReturns(name)
+	return t
+}
diff --git a/pkg/rtc/room.go b/pkg/rtc/room.go
index 310b544f7..b10d707ad 100644
--- a/pkg/rtc/room.go
+++ b/pkg/rtc/room.go
@@ -144,6 +144,10 @@ func (r *Room) onTrackAdded(participant Participant, track PublishedTrack) {
 			// skip publishing peer
 			continue
 		}
+		if existingParticipant.State() != livekit.ParticipantInfo_JOINED {
+			// not fully joined. don't subscribe yet
+			continue
+		}
 		if err := track.AddSubscriber(existingParticipant); err != nil {
 			logger.GetLogger().Errorw("could not subscribe to remoteTrack",
 				"srcParticipant", participant.ID(),
diff --git a/pkg/rtc/room_test.go b/pkg/rtc/room_test.go
index 416a39dd5..e641e7fa4 100644
--- a/pkg/rtc/room_test.go
+++ b/pkg/rtc/room_test.go
@@ -2,6 +2,7 @@ package rtc_test
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 
@@ -10,18 +11,21 @@ import (
 	"github.com/livekit/livekit-server/proto/livekit"
 )
 
+const (
+	numParticipants = 3
+)
+
 func TestRoomJoin(t *testing.T) {
 	t.Run("joining returns existing participant data", func(t *testing.T) {
-		numExisting := 3
-		rm := newRoomWithParticipants(t, numExisting)
+		rm := newRoomWithParticipants(t, numParticipants)
 		p := newMockParticipant("new")
 
 		rm.Join(p)
 
 		// expect new participant to get a JoinReply
 		participants := p.SendJoinResponseArgsForCall(0)
-		assert.Len(t, participants, numExisting)
-		assert.Len(t, rm.GetParticipants(), numExisting+1)
+		assert.Len(t, participants, numParticipants)
+		assert.Len(t, rm.GetParticipants(), numParticipants+1)
 	})
 
 	t.Run("subscribe to existing channels upon join", func(t *testing.T) {
@@ -48,6 +52,47 @@ func TestRoomJoin(t *testing.T) {
 			assert.Equal(t, p, mockP.AddSubscriberArgsForCall(mockP.AddSubscriberCallCount()-1))
 		}
 	})
+
+	t.Run("participant removal is broadcasted to others", func(t *testing.T) {
+		rm := newRoomWithParticipants(t, numParticipants)
+		participants := rm.GetParticipants()
+		p := participants[0].(*rtcfakes.FakeParticipant)
+
+		rm.RemoveParticipant(p.ID())
+		time.Sleep(10 * time.Millisecond)
+
+		for _, op := range participants {
+			if op == p {
+				assert.Zero(t, p.SendParticipantUpdateCallCount())
+				continue
+			}
+			fakeP := op.(*rtcfakes.FakeParticipant)
+			assert.Equal(t, 1, fakeP.SendParticipantUpdateCallCount())
+		}
+	})
+}
+
+func TestNewTrack(t *testing.T) {
+	t.Run("new track should be added to connected participants", func(t *testing.T) {
+		rm := newRoomWithParticipants(t, 4)
+		participants := rm.GetParticipants()
+		p0 := participants[0].(*rtcfakes.FakeParticipant)
+		p0.StateReturns(livekit.ParticipantInfo_JOINING)
+		p1 := participants[1].(*rtcfakes.FakeParticipant)
+		p1.StateReturns(livekit.ParticipantInfo_DISCONNECTED)
+		p2 := participants[2].(*rtcfakes.FakeParticipant)
+		p2.StateReturns(livekit.ParticipantInfo_JOINED)
+		p3 := participants[3].(*rtcfakes.FakeParticipant)
+
+		// p3 adds track
+		track := newMockTrack(livekit.TrackInfo_VIDEO, "webcam")
+		trackCB := p3.OnTrackPublishedArgsForCall(0)
+		assert.NotNil(t, trackCB)
+		trackCB(p3, track)
+		// only p2 should've been called
+		assert.Equal(t, 1, track.AddSubscriberCallCount())
+		assert.Equal(t, p2, track.AddSubscriberArgsForCall(0))
+	})
 }
 
 func newRoomWithParticipants(t *testing.T, num int) *rtc.Room {