import org.junit.Test;

import java.util.LinkedList;
import java.util.Queue;

import static org.junit.Assert.*;

public class UserGraphTest {

    private UserGraph generateTestUserGraph() {
        User coach1 = new User.Builder().build("Coach1");
        User coach2 = new User.Builder().build("Coach2");

        User student1 = new User.Builder().addCoach(coach1).build("Student1");
        User student2 = new User.Builder().addCoach(coach1).build("Student2");
        User student3 = new User.Builder().addCoach(coach1).addCoach(coach2).build("Student3");
        User student4 = new User.Builder().addCoach(coach2).build("Student4");
        User student5 = new User.Builder().addCoach(coach2).build("Student5");

        return new UserGraph.Builder()
                .addUser(coach1)
                .addUser(coach2)
                .addUser(student1)
                .addUser(student2)
                .addUser(student3)
                .addUser(student4)
                .addUser(student5)
                .build();
    }

    @Test
    public void testInfectSingleUser() {
        User user1 = new User.Builder().build("TestUser");
        Queue<User> infectQueue = new LinkedList<>();
        Version newVersion = new Version(1, 1, 1);

        UserGraph testGraph = new UserGraph.Builder().addUser(user1).build();
        testGraph.infectUser(user1, infectQueue, newVersion);

        assertEquals("Version updated correctly", new Version(1, 1, 1), user1.getCurrentVersion());
        assertEquals("No users added to queue", new LinkedList<User>(), infectQueue);
    }

    @Test
    public void testInfectSeveralUsers() {
        UserGraph testGraph = generateTestUserGraph();
        testGraph.getUser("Student5").addStudent(testGraph.getUser("Student1"));


        Queue<User> infectQueue = new LinkedList<>();
        Version newVersion = new Version(1, 1, 1);

        testGraph.infectUser(testGraph.getUser("Student5"), infectQueue, newVersion);

        Queue<User> expectedQueue = new LinkedList<>();
        expectedQueue.add(testGraph.getUser("Coach2"));
        expectedQueue.add(testGraph.getUser("Student1"));

        assertEquals("Version updated correctly", new Version(1, 1, 1), testGraph.getUser("Student5").getCurrentVersion());
        assertEquals("No users added to queue", expectedQueue, infectQueue);
    }
}
@@ -712,16 +712,78 @@
 dent5%22);
+%0A        User student6 = new User.Builder().build(%22Student6%22);
 %0A%0A      
@@ -1056,16 +1056,51 @@
 udent5)%0A
+                .addUser(student6)%0A
         
@@ -1864,17 +1864,16 @@
 t1%22));%0A%0A
-%0A
         
@@ -2344,32 +2344,32 @@
 rentVersion());%0A
-
         assertEq
@@ -2434,10 +2434,682 @@
 );%0A    %7D
+%0A%0A    @Test%0A    public void testTotalInfection() %7B%0A        UserGraph testGraph = generateTestUserGraph();%0A        Version newVersion = new Version(1, 1, 1);%0A        testGraph.totalInfection(testGraph.getUser(%22Student1%22), newVersion);%0A        for (String username : testGraph.getUsers().keySet()) %7B%0A            User user = testGraph.getUser(username);%0A            if (!username.equals(%22Student6%22)) %7B%0A                assertEquals(%22User not infected%22, new Version(1, 1, 1), user.getCurrentVersion());%0A            %7D%0A            else %7B%0A                assertEquals(%22User infected but shouldn't be%22, new Version(1, 0, 0), user.getCurrentVersion());%0A            %7D%0A        %7D%0A    %7D
 %0A%7D

Test for total infection.