diff --git a/lom/src/main/kotlin/lom.gradle.kts b/lom/src/main/kotlin/lom.gradle.kts
index 7bda193..c92c670 100644
--- a/lom/src/main/kotlin/lom.gradle.kts
+++ b/lom/src/main/kotlin/lom.gradle.kts
@@ -29,15 +29,20 @@ loom {
 // Create testmod source set with access to main and client classpaths
 sourceSets {
     create("testmod") {
-        compileClasspath += sourceSets.main.get().compileClasspath
-        runtimeClasspath += sourceSets.main.get().runtimeClasspath
+        compileClasspath += main.get().compileClasspath
+        runtimeClasspath += main.get().runtimeClasspath
 
         if (args.isSplitSources) {
-            compileClasspath += sourceSets.client.get().compileClasspath
-            runtimeClasspath += sourceSets.client.get().runtimeClasspath
+            compileClasspath += client.get().compileClasspath
+            runtimeClasspath += client.get().runtimeClasspath
         }
     }
-    if (!args.isSplitSources) {
+    if (args.isSplitSources) {
+        test {
+            compileClasspath += client.get().compileClasspath
+            runtimeClasspath += client.get().runtimeClasspath
+        }
+    } else {
         main {
             java.srcDir("src/client/java")
             resources.srcDir("src/client/resources")