diff --git a/kunit.sh b/kunit.sh
index 4116104..14d0319 100755
--- a/kunit.sh
+++ b/kunit.sh
@@ -26,7 +26,12 @@
 echo -e "CONFIG_DEBUG_KERNEL=y\nCONFIG_DEBUG_INFO=y\nCONFIG_GCOV=y" >> .config
 
 # Run the KUnit tool unit tests
-python3 ./tools/testing/kunit/kunit_test.py > "$ARTIFACTS/kunit_test.log" 2>&1
+if [[ -f ./tools/testing/kunit/kunit_test.py ]]
+  python3 ./tools/testing/kunit/kunit_test.py > "$ARTIFACTS/kunit_test.log" 2>&1
+else # for upstream version
+  python3 ./tools/testing/kunit/kunit_tool_test.py > "$ARTIFACTS/kunit_test.log" 2>&1
+fi
+
 KUNIT_TEST_EXIT_CODE=$?
 
 # Build and run kernel, parse output, and send details to output log