diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgMethodConstructor.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgMethodConstructor.kt
index c9b7f2f231..85d28ac4a2 100644
--- a/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgMethodConstructor.kt
+++ b/utbot-framework/src/main/kotlin/org/utbot/framework/codegen/model/constructor/tree/CgMethodConstructor.kt
@@ -1004,18 +1004,17 @@ internal class CgMethodConstructor(val context: CgContext) : CgContextOwner by c
                         } else {
                             // array of objects, have to use deep equals
 
-                            if (expected is CgLiteral) {
-                                // Literal can only be Primitive or String, can use equals here
-                                testFrameworkManager.assertEquals(expected, actual)
-                                return
-                            }
-
-                            require(resultModel is UtArrayModel) {
-                                "Result model have to be UtArrayModel to generate arrays assertion " +
-                                        "but `${resultModel::class}` found"
+                            when (expected) {
+                                is CgLiteral -> testFrameworkManager.assertEquals(expected, actual)
+                                is CgNotNullAssertion -> generateForNotNullAssertion(expected, actual)
+                                else -> {
+                                    require(resultModel is UtArrayModel) {
+                                        "Result model have to be UtArrayModel to generate arrays assertion " +
+                                                "but `${resultModel::class}` found"
+                                    }
+                                    generateDeepEqualsOrNullAssertion(expected, actual)
+                                }
                             }
-
-                            generateDeepEqualsOrNullAssertion(expected, actual)
                         }
                     }
                 }
@@ -1046,16 +1045,8 @@ internal class CgMethodConstructor(val context: CgContext) : CgContextOwner by c
                 }
                 else -> {
                     when (expected) {
-                        is CgLiteral -> {
-                            // Literal can only be Primitive or String, can use equals here
-                            testFrameworkManager.assertEquals(expected, actual)
-                        }
-                        is CgNotNullAssertion -> {
-                            require(expected.expression is CgVariable) {
-                                "Only Variable wrapped in CgNotNullAssertion is supported in deep equals"
-                            }
-                            currentBlock = currentBlock.addAll(generateDeepEqualsAssertion(expected.expression, actual))
-                        }
+                        is CgLiteral -> testFrameworkManager.assertEquals(expected, actual)
+                        is CgNotNullAssertion -> generateForNotNullAssertion(expected, actual)
                         else -> generateDeepEqualsOrNullAssertion(expected, actual)
                     }
                 }
@@ -1063,6 +1054,13 @@ internal class CgMethodConstructor(val context: CgContext) : CgContextOwner by c
         }
     }
 
+    private fun generateForNotNullAssertion(expected: CgNotNullAssertion, actual: CgVariable) {
+        require(expected.expression is CgVariable) {
+            "Only CgVariable wrapped in CgNotNullAssertion is supported in deepEquals"
+        }
+        generateDeepEqualsOrNullAssertion(expected.expression, actual)
+    }
+
     /**
      * We can't use standard deepEquals method in parametrized tests
      * because nullable objects require different asserts.