diff --git a/build.gradle.kts b/build.gradle.kts
index 51e935e..03ca216 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -8,7 +8,7 @@ plugins {
 }
 
 group = "org.jetbrains.academy.test.system"
-version = "2.1.1"
+version = "2.1.2"
 
 allprojects {
     apply {
diff --git a/core/src/main/kotlin/org/jetbrains/academy/test/system/core/MethodUtils.kt b/core/src/main/kotlin/org/jetbrains/academy/test/system/core/MethodUtils.kt
index d411a5b..bc7b532 100644
--- a/core/src/main/kotlin/org/jetbrains/academy/test/system/core/MethodUtils.kt
+++ b/core/src/main/kotlin/org/jetbrains/academy/test/system/core/MethodUtils.kt
@@ -39,9 +39,9 @@ fun Array<Method>.findMethod(method: TestMethod, customErrorMessage: String? = n
                 it.name == method.name
             }
         }
-    val returnTypeJava = (listOf(method.returnTypeJava)).map { it.lowercase() }
+    val returnTypeJava = (method.returnTypeJava?.let { listOf(it) } ?: listOfNotNull(method.returnType.type, *method.returnType.possibleBounds.toTypedArray())).map { it.lowercase() }
     val filteredByType =
-        filteredByName.filterByCondition(customErrorMessage ?: "The method ${method.name} should have the return type ${method.returnType?.getTypePrettyString()}") {
+        filteredByName.filterByCondition(customErrorMessage ?: "The method ${method.name} should have the return type ${method.returnType.getTypePrettyString()}") {
             it.returnType.name.getShortName().lowercase() in returnTypeJava
         }
     val filteredByArgumentsCount =
diff --git a/core/src/main/kotlin/org/jetbrains/academy/test/system/core/models/method/TestMethod.kt b/core/src/main/kotlin/org/jetbrains/academy/test/system/core/models/method/TestMethod.kt
index 45c49ce..34a7f2b 100644
--- a/core/src/main/kotlin/org/jetbrains/academy/test/system/core/models/method/TestMethod.kt
+++ b/core/src/main/kotlin/org/jetbrains/academy/test/system/core/models/method/TestMethod.kt
@@ -20,9 +20,9 @@ import kotlin.reflect.jvm.kotlinFunction
  */
 data class TestMethod(
     val name: String,
-    val returnTypeJava: String,
-    val returnType: TestKotlinType? = null,
+    val returnType: TestKotlinType,
     val arguments: List<TestVariable> = emptyList(),
+    val returnTypeJava: String? = null,
     val visibility: Visibility = Visibility.PUBLIC,
     val hasGeneratedPartInName: Boolean = false,
 ) {
@@ -33,15 +33,13 @@ data class TestMethod(
         visibility: Visibility
     ) : this(
         name = name,
-        returnTypeJava = returnTypeJava,
-        returnType = null,
+        returnType = TestKotlinType(returnTypeJava),
         arguments = arguments,
+        returnTypeJava = returnTypeJava,
         visibility = visibility,
         hasGeneratedPartInName = false
     )
 
-    private fun getTypePrettyString() = returnType?.getTypePrettyString() ?: returnTypeJava
-
     fun prettyString(withToDo: Boolean = true): String {
         val args = arguments.joinToString(", ") { it.paramPrettyString() }
         val body = if (withToDo) {
@@ -49,7 +47,7 @@ data class TestMethod(
         } else {
             "// Some code"
         }
-        return "${visibility.key} fun $name($args): ${getTypePrettyString()} = $body"
+        return "${visibility.key} fun $name($args): ${returnType.getTypePrettyString()} = $body"
     }
 
     private fun TestVariable.paramPrettyString() = "$name: $javaType"
@@ -64,6 +62,6 @@ data class TestMethod(
             this.visibility.key,
             "\"The visibility of the method $name must be ${this.visibility.key}\""
         )
-        kotlinFunction.returnType.checkType(returnType, returnTypeJava, "the function $name")
+        kotlinFunction.returnType.checkType(returnType, returnTypeJava ?: returnType.type, "the function $name")
     }
 }
diff --git a/core/src/test/kotlin/org/jetbrains/academy/test/system/core/testData/java/TestJavaClass.kt b/core/src/test/kotlin/org/jetbrains/academy/test/system/core/testData/java/TestJavaClass.kt
index 5e458d1..0d4fc7d 100644
--- a/core/src/test/kotlin/org/jetbrains/academy/test/system/core/testData/java/TestJavaClass.kt
+++ b/core/src/test/kotlin/org/jetbrains/academy/test/system/core/testData/java/TestJavaClass.kt
@@ -81,11 +81,13 @@ val javaClassTestClass = TestClass(
         TestMethod(
             name = "publicMethod",
             returnTypeJava = "void",
+            arguments = emptyList(),
             visibility = Visibility.PUBLIC
         ),
         TestMethod(
             name = "privateMethod",
             returnTypeJava = "void",
+            arguments = emptyList(),
             visibility = Visibility.PRIVATE
         ),
         TestMethod(