diff --git a/idea/resources/inspectionDescriptions/MainFunctionReturnUnit.html b/idea/resources/inspectionDescriptions/MainFunctionReturnUnit.html
index 7043219f03d..026b0ead172 100644
--- a/idea/resources/inspectionDescriptions/MainFunctionReturnUnit.html
+++ b/idea/resources/inspectionDescriptions/MainFunctionReturnUnit.html
@@ -1,5 +1,5 @@
 <html>
 <body>
-This inspection reports a main function with incorrect return type (should be <b>Unit</b>).
+This inspection reports a entry point function with incorrect return type (should be <b>Unit</b>).
 </body>
 </html>
diff --git a/idea/src/META-INF/plugin.xml b/idea/src/META-INF/plugin.xml
index 5fb736da19f..ad7edcab688 100644
--- a/idea/src/META-INF/plugin.xml
+++ b/idea/src/META-INF/plugin.xml
@@ -2856,7 +2856,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/META-INF/plugin.xml.172 b/idea/src/META-INF/plugin.xml.172
index 5a22051e866..35af67d21bb 100644
--- a/idea/src/META-INF/plugin.xml.172
+++ b/idea/src/META-INF/plugin.xml.172
@@ -2854,7 +2854,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/META-INF/plugin.xml.173 b/idea/src/META-INF/plugin.xml.173
index 15a7d7f6b7b..0bda9b150ec 100644
--- a/idea/src/META-INF/plugin.xml.173
+++ b/idea/src/META-INF/plugin.xml.173
@@ -2854,7 +2854,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/META-INF/plugin.xml.182 b/idea/src/META-INF/plugin.xml.182
index 53f894c0453..533bc8680e7 100644
--- a/idea/src/META-INF/plugin.xml.182
+++ b/idea/src/META-INF/plugin.xml.182
@@ -2855,7 +2855,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/META-INF/plugin.xml.as31 b/idea/src/META-INF/plugin.xml.as31
index d516619c07e..92f76a2c485 100644
--- a/idea/src/META-INF/plugin.xml.as31
+++ b/idea/src/META-INF/plugin.xml.as31
@@ -2854,7 +2854,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/META-INF/plugin.xml.as32 b/idea/src/META-INF/plugin.xml.as32
index e5c9c867105..f69651737f0 100644
--- a/idea/src/META-INF/plugin.xml.as32
+++ b/idea/src/META-INF/plugin.xml.as32
@@ -2854,7 +2854,7 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
     />
 
     <localInspection implementationClass="org.jetbrains.kotlin.idea.inspections.MainFunctionReturnUnitInspection"
-                     displayName="Main function should return Unit"
+                     displayName="Entry point function should return Unit"
                      groupPath="Kotlin"
                      groupName="Probable bugs"
                      enabledByDefault="true"
diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt
index ae5f3b0bb2a..651e1b04177 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/MainFunctionReturnUnitInspection.kt
@@ -10,32 +10,41 @@ import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.MainFunctionDetector
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.setType
 import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.namedFunctionVisitor
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 class MainFunctionReturnUnitInspection : AbstractKotlinInspection() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
         return namedFunctionVisitor(fun(function: KtNamedFunction) {
-            if (function.name != "main") return
+            val isMain = function.name == "main"
+            val testAnnotations = function.annotationEntries.filter { "Test" in it.text }
+            if (!isMain && testAnnotations.isEmpty()) return
 
             val descriptor = function.descriptor as? FunctionDescriptor ?: return
-            if (!MainFunctionDetector.isMain(descriptor, checkReturnType = false)) return
-            if (MainFunctionDetector.isMainReturnType(descriptor)) return
+            if (isMain) {
+                if (!MainFunctionDetector.isMain(descriptor, checkReturnType = false)) return
+            } else {
+                val junitTestFqNames = listOf(FqName("org.junit.Test"), FqName("org.junit.jupiter.api.Test"))
+                if (testAnnotations.none { it.fqName() in junitTestFqNames }) return
+            }
+            if (descriptor.returnType?.let { KotlinBuiltIns.isUnit(it) } == true) return
 
             holder.registerProblem(
                 function.nameIdentifier ?: function,
-                "main should return Unit",
+                "${if (isMain) "main" else "JUnit test"} should return Unit",
                 ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
                 ChangeMainFunctionReturnTypeToUnitFix(function.typeReference != null)
             )
@@ -43,6 +52,12 @@ class MainFunctionReturnUnitInspection : AbstractKotlinInspection() {
     }
 }
 
+private fun KtAnnotationEntry.fqName(): FqName? {
+    val typeReference = this.typeReference ?: return null
+    val context = typeReference.analyze(BodyResolveMode.PARTIAL)
+    return context[BindingContext.TYPE, typeReference]?.constructor?.declarationDescriptor?.fqNameSafe
+}
+
 private class ChangeMainFunctionReturnTypeToUnitFix(private val hasExplicitReturnType: Boolean) : LocalQuickFix {
     override fun getName() = if (hasExplicitReturnType) "Change return type to Unit" else "Add explicit Unit return type"
 
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt
new file mode 100644
index 00000000000..e88385ce0be
--- /dev/null
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt
@@ -0,0 +1,10 @@
+// PROBLEM: JUnit test should return Unit
+// FIX: Add explicit Unit return type
+// DISABLE-ERRORS
+package org.junit
+
+annotation class Test
+
+class A {
+    @Test fun foo<caret>() = 1
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt.after
new file mode 100644
index 00000000000..3e908b6c8cc
--- /dev/null
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt.after
@@ -0,0 +1,10 @@
+// PROBLEM: JUnit test should return Unit
+// FIX: Add explicit Unit return type
+// DISABLE-ERRORS
+package org.junit
+
+annotation class Test
+
+class A {
+    @Test fun foo<caret>(): Unit = 1
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt
new file mode 100644
index 00000000000..d3f4eddeeb0
--- /dev/null
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt
@@ -0,0 +1,9 @@
+// FIX: Add explicit Unit return type
+// DISABLE-ERRORS
+package org.junit.jupiter.api
+
+annotation class Test
+
+class A {
+    @Test fun foo<caret>() = 1
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt.after
new file mode 100644
index 00000000000..ae475e95a63
--- /dev/null
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt.after
@@ -0,0 +1,9 @@
+// FIX: Add explicit Unit return type
+// DISABLE-ERRORS
+package org.junit.jupiter.api
+
+annotation class Test
+
+class A {
+    @Test fun foo<caret>(): Unit = 1
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt
index f5effba3b7a..de1a6a306e6 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt
@@ -1,3 +1,4 @@
+// PROBLEM: main should return Unit
 // FIX: Add explicit Unit return type
 // DISABLE-ERRORS
 fun <caret>main(args: Array<String>) = 1
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after
index f3bb7eeaee2..a508b572c4a 100644
--- a/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after
+++ b/idea/testData/inspectionsLocal/mainFunctionReturnUnit/notUnit1.kt.after
@@ -1,3 +1,4 @@
+// PROBLEM: main should return Unit
 // FIX: Add explicit Unit return type
 // DISABLE-ERRORS
 fun main(args: Array<String>): Unit = 1
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
index 88ea97f6000..90b00ea0e9e 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java
@@ -2383,6 +2383,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/inspectionsLocal/mainFunctionReturnUnit"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("junit4Test.kt")
+        public void testJunit4Test() throws Exception {
+            runTest("idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit4Test.kt");
+        }
+
+        @TestMetadata("junit5Test.kt")
+        public void testJunit5Test() throws Exception {
+            runTest("idea/testData/inspectionsLocal/mainFunctionReturnUnit/junit5Test.kt");
+        }
+
         @TestMetadata("notMain.kt")
         public void testNotMain() throws Exception {
             runTest("idea/testData/inspectionsLocal/mainFunctionReturnUnit/notMain.kt");
