diff --git a/build.gradle b/build.gradle
index 2b52d53cd58a1a20b24d249d409b5d274222f9d2..a3c39ae7078adbb140b6484c6cbbb89263ef3c02 100644
--- a/build.gradle
+++ b/build.gradle
@@ -25,6 +25,7 @@ dependencies {
 
 test {
     useJUnitPlatform()
+    maxHeapSize = "4096m"
 }
 
 java {
diff --git a/src/main/java/chat/UMessageServerConnection.java b/src/main/java/chat/UMessageServerConnection.java
index a5fd5ef01503321d8515d40942dfa5211bb272f3..e810a9a8e09ccfc405d82932119ff88ea68f4c6b 100644
--- a/src/main/java/chat/UMessageServerConnection.java
+++ b/src/main/java/chat/UMessageServerConnection.java
@@ -1,5 +1,6 @@
 package chat;
 
+import javax.swing.*;
 import java.io.*;
 import java.net.Socket;
 import java.net.UnknownHostException;
@@ -72,12 +73,12 @@ public class UMessageServerConnection extends Thread {
         try {
             cmd = cmd.trim();
             switch (cmd) {
-            case "MAIN":
-                m_channel(UMessageServerConnection.CHAT_CHANNEL, text);
-                break;
-            default:
-                String[] parts = text.split(" ", 2);
-                cmd = parts[0].trim();
+                case "MAIN":
+                    m_channel(UMessageServerConnection.CHAT_CHANNEL, text);
+                    break;
+                default:
+                    String[] parts = text.split(" ", 2);
+                    cmd = parts[0].trim();
             }
         } catch (Exception e) {
             e.printStackTrace();
@@ -112,34 +113,52 @@ public class UMessageServerConnection extends Thread {
             while ((line = this.in.readLine()) != null) {
                 if (line.startsWith("PING")) {
                     write("PONG", line.substring(5));
-                }
-                else if (line.startsWith(":umessage")) {
+                } else if (line.split(" ")[1].trim().equals("353")) {
                     int code = Integer.parseInt(line.split(" ")[1]);
                     switch (code) {
-                    // List of users...
-                    case IRCCodes.RplNamReply:
-                        String[] names = line.split(":")[2].split(" ");
-                        this.main.window.addUsers(names);
-                        break;
+                        // List of users...
+                        case IRCCodes.RplNamReply:
+                            String[] names = line.split(":")[2].split(" ");
+                            SwingUtilities.invokeLater(new Runnable() {
+                                @Override
+                                public void run() {
+                                    main.window.addUsers(names);
+                                }
+                            });
+                            break;
                     }
-                }
-                else {
+                } else {
                     String cmd = line.split(" ")[1];
+                    final String[] lineParts = line.split(":");
                     switch (cmd) {
-                    case "PART":
-                        this.main.window.removeUser(line.split(":")[1].split("!")[0]);
-                        break;
-                    case "JOIN":
-                        this.main.window.addUser(line.split(":")[1].split("!")[0]);
-                        break;
-                    case "PRIVMSG":
-                        if (!line.split(" ")[2].equals(this.username)) {
-                            return;
-                        }
-                        String[] lineParts = line.split(":");
-                        this.main.window.gotMessage(lineParts[1].split("!")[0],
-                                lineParts[2]);
-                        break;
+                        case "PART":
+                            SwingUtilities.invokeLater(new Runnable() {
+                                @Override
+                                public void run() {
+                                    main.window.removeUser(lineParts[1].split("!")[0]);
+                                }
+                            });
+                            break;
+                        case "JOIN":
+                            SwingUtilities.invokeLater(new Runnable() {
+                                @Override
+                                public void run() {
+                                    main.window.addUser(lineParts[1].split("!")[0]);
+                                }
+                            });
+                            break;
+                        case "PRIVMSG":
+                            if (!line.split(" ")[2].equals(this.username)) {
+                                return;
+                            }
+                            SwingUtilities.invokeLater(new Runnable() {
+                                @Override
+                                public void run() {
+                                    main.window.gotMessage(lineParts[1].split("!")[0],
+                                            lineParts[2]);
+                                }
+                            });
+                            break;
                     }
                 }
             }
diff --git a/src/main/java/datastructures/dictionaries/AVLTree.java b/src/main/java/datastructures/dictionaries/AVLTree.java
index 8677d00a8459ceb3d1515c8a957f389a122de3db..dce422bb4d0133f31ffb32502b2176f22ace1d81 100644
--- a/src/main/java/datastructures/dictionaries/AVLTree.java
+++ b/src/main/java/datastructures/dictionaries/AVLTree.java
@@ -19,9 +19,9 @@ import cse332.datastructures.trees.BinarySearchTree;
  * lead to highly perplexing and erroneous behavior. Instead,
  * continue using the existing BSTNode children array.
  * 4. Ensure that the class does not have redundant methods
- * 5. Cast children array to AVLNode whenever necessary in your
- * AVLTree. This will result a lot of casts, so we recommend you make
- * private methods that encapsulate those casts.
+ * 5. Cast a BSTNode to an AVLNode whenever necessary in your AVLTree.
+ * This will result a lot of casts, so we recommend you make private methods
+ * that encapsulate those casts.
  * 6. Do NOT override the toString method. It is used for grading.
  * 7. The internal structure of your AVLTree (from this.root to the leaves) must be correct
  */
diff --git a/src/test/java/ckpt2/TopKSortTests.java b/src/test/java/ckpt2/TopKSortTests.java
index f0a52f3a078f63577aac8e8c3e14581c64d02d2f..741adf1166a3132385de088e25c2634acdf65860 100644
--- a/src/test/java/ckpt2/TopKSortTests.java
+++ b/src/test/java/ckpt2/TopKSortTests.java
@@ -18,7 +18,7 @@ public class TopKSortTests {
 		Integer[] arr_sorted = {7, 8, 9, 10};
 		TopKSort.sort(arr, K, Integer::compareTo);
 		for(int i = 0; i < K; i++) {
-			assertEquals(arr[i], arr_sorted[i]);
+			assertEquals(arr_sorted[i], arr[i]);
 		}
 	}
 
@@ -30,7 +30,7 @@ public class TopKSortTests {
 		Integer[] arr_sorted = {6, 7, 8, 9};
 		TopKSort.sort(arr, K, Integer::compareTo);
 		for(int i = 0; i < K; i++) {
-			assertEquals(arr[i], arr_sorted[i]);
+			assertEquals(arr_sorted[i], arr[i]);
 		}
 	}
-}
\ No newline at end of file
+}