diff --git a/src/Generator/Helper/CodeHelper.php b/src/Generator/Helper/CodeHelper.php
index 0bd1389..cb96a72 100644
--- a/src/Generator/Helper/CodeHelper.php
+++ b/src/Generator/Helper/CodeHelper.php
@@ -524,7 +524,6 @@ public function parseMethodsFromSource(string $sourceName): array
      */
     protected function getMethodReturnType(ReflectionMethod $reflectionMethod): string
     {
-        $returnType = null;
         $className = $reflectionMethod->getDeclaringClass()->getName();
 
         if (in_array($reflectionMethod->getName(), self::METHOD_NAMES_RETURN_SELF, true)) {
@@ -534,7 +533,7 @@ protected function getMethodReturnType(ReflectionMethod $reflectionMethod): stri
 
             if (
                 class_exists('ReflectionUnionType') &&
-                $className instanceof \ReflectionUnionType
+                $returnType instanceof \ReflectionUnionType
             ) {
                 // use only first return type
                 // @TODO return all return types
@@ -546,7 +545,7 @@ class_exists('ReflectionUnionType') &&
             }
         }
 
-        if (!$returnType || $returnType === 'array') {
+        if (!$returnType || $returnType === DataTypeInterface::TYPE_ARRAY) {
             $returnType = $this->getReturnFromReflectionMethodAnnotation($reflectionMethod, false, true);
 
             if (null === $returnType) {
diff --git a/src/Generator/Mock/Mockery.php b/src/Generator/Mock/Mockery.php
index 5495f15..e549b54 100644
--- a/src/Generator/Mock/Mockery.php
+++ b/src/Generator/Mock/Mockery.php
@@ -117,7 +117,7 @@ public function generateMockCode(
 
                 if (
                     class_exists('ReflectionUnionType') &&
-                    $className instanceof \ReflectionUnionType
+                    $returnType instanceof \ReflectionUnionType
                 ) {
                     // use only first return type
                     // @TODO return all return types