diff --git a/src/bitExpert/PHPStan/Magento/Reflection/AbstractMagicMethodReflectionExtension.php b/src/bitExpert/PHPStan/Magento/Reflection/AbstractMagicMethodReflectionExtension.php
index 8650066..8571456 100644
--- a/src/bitExpert/PHPStan/Magento/Reflection/AbstractMagicMethodReflectionExtension.php
+++ b/src/bitExpert/PHPStan/Magento/Reflection/AbstractMagicMethodReflectionExtension.php
@@ -17,9 +17,15 @@
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\Reflection\MethodsClassReflectionExtension;
 use PHPStan\Reflection\Php\DummyParameter;
-use PHPStan\Reflection\TrivialParametersAcceptor;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\BooleanType;
 use PHPStan\Type\Generic\TemplateTypeMap;
+use PHPStan\Type\IntegerType;
 use PHPStan\Type\MixedType;
+use PHPStan\Type\NullType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\StringType;
+use PHPStan\Type\UnionType;
 
 abstract class AbstractMagicMethodReflectionExtension implements MethodsClassReflectionExtension
 {
@@ -30,47 +36,143 @@ abstract class AbstractMagicMethodReflectionExtension implements MethodsClassRef
      */
     public function getMethod(ClassReflection $classReflection, string $methodName): MethodReflection
     {
-        if (strpos($methodName, 'get') === 0) {
-            return $this->returnGetMagicMethodReflection($classReflection, $methodName);
+        $methodPrefix = substr($methodName, 0, 3);
+        switch ($methodPrefix) {
+            case 'get':
+                return $this->returnGetMagicMethod($classReflection, $methodName);
+            case 'set':
+                return $this->returnSetMagicMethod($classReflection, $methodName);
+            case 'uns':
+                return $this->returnUnsetMagicMethod($classReflection, $methodName);
+            case 'has':
+                return $this->returnHasMagicMethod($classReflection, $methodName);
+            default:
+                break;
         }
-
-        return $this->returnMagicMethodReflection($classReflection, $methodName);
     }
 
     /**
-     * Helper method to create magic method for get calls. The method call does not accept any parameter and will return
-     * mixed type.
+     * Helper method to create magic method reflection for get() calls.
      *
      * @param ClassReflection $classReflection
      * @param string $methodName
      * @return MethodReflection
      */
-    protected function returnGetMagicMethodReflection(
-        ClassReflection $classReflection,
-        string $methodName
-    ): MethodReflection {
-        $variants = new TrivialParametersAcceptor();
+    private function returnGetMagicMethod(ClassReflection $classReflection, string $methodName): MethodReflection
+    {
+        $params = [
+            new DummyParameter(
+                'key',
+                new StringType(),
+                true,
+                null,
+                false,
+                null
+            ),
+            new DummyParameter(
+                'index',
+                new UnionType([new StringType(), new IntegerType(), new NullType()]),
+                true,
+                null,
+                false,
+                null
+            )
+        ];
+
+        $returnType = new MixedType();
+
+        $variants = new FunctionVariant(
+            TemplateTypeMap::createEmpty(),
+            null,
+            $params,
+            false,
+            $returnType
+        );
+
         return new MagicMethodReflection($methodName, $classReflection, [$variants]);
     }
 
-    /**
-     * Helper method to create magic reflection method for set, unset and has calls. Those method calls accept one
-     * mixed parameter and return mixed type.
-     *
-     * @param ClassReflection $classReflection
-     * @param string $methodName
-     * @return MethodReflection
-     */
-    protected function returnMagicMethodReflection(
-        ClassReflection $classReflection,
-        string $methodName
-    ): MethodReflection {
+    private function returnSetMagicMethod(ClassReflection $classReflection, string $methodName): MethodReflection
+    {
+        $params = [
+            new DummyParameter(
+                'key',
+                new UnionType([new StringType(), new ArrayType(new MixedType(), new MixedType())]),
+                true,
+                null,
+                false,
+                null
+            ),
+            new DummyParameter(
+                'value',
+                new MixedType(),
+                true,
+                null,
+                false,
+                null
+            )
+        ];
+
+        $returnType = new ObjectType($classReflection->getName());
+
+        $variants = new FunctionVariant(
+            TemplateTypeMap::createEmpty(),
+            null,
+            $params,
+            false,
+            $returnType
+        );
+
+        return new MagicMethodReflection($methodName, $classReflection, [$variants]);
+    }
+
+    private function returnUnsetMagicMethod(ClassReflection $classReflection, string $methodName): MethodReflection
+    {
+        $params = [
+            new DummyParameter(
+                'key',
+                new UnionType([new NullType(), new StringType(), new ArrayType(new MixedType(), new MixedType())]),
+                true,
+                null,
+                false,
+                null
+            )
+        ];
+
+        $returnType = new ObjectType($classReflection->getName());
+
+        $variants = new FunctionVariant(
+            TemplateTypeMap::createEmpty(),
+            null,
+            $params,
+            false,
+            $returnType
+        );
+
+        return new MagicMethodReflection($methodName, $classReflection, [$variants]);
+    }
+
+    private function returnHasMagicMethod(ClassReflection $classReflection, string $methodName): MethodReflection
+    {
+        $params = [
+            new DummyParameter(
+                'key',
+                new StringType(),
+                true,
+                null,
+                false,
+                null
+            )
+        ];
+
+        $returnType = new BooleanType();
+
         $variants = new FunctionVariant(
             TemplateTypeMap::createEmpty(),
             null,
-            [ new DummyParameter('name', new MixedType(), false, null, false, null),],
+            $params,
             false,
-            new MixedType()
+            $returnType
         );
 
         return new MagicMethodReflection($methodName, $classReflection, [$variants]);
diff --git a/src/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtension.php b/src/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtension.php
index 7caea8a..32b4077 100644
--- a/src/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtension.php
+++ b/src/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtension.php
@@ -26,11 +26,9 @@ class DataObjectMagicMethodReflectionExtension extends AbstractMagicMethodReflec
     public function hasMethod(ClassReflection $classReflection, string $methodName): bool
     {
         $parentClasses = $classReflection->getParentClassesNames();
+        $parentClasses[] = $classReflection->getName();
 
-        $isDataObject = $classReflection->getName() === DataObject::class ||
-                        in_array(DataObject::class, $parentClasses, true);
-
-        return $isDataObject &&
+        return in_array(DataObject::class, $parentClasses, true) &&
             in_array(substr($methodName, 0, 3), ['get', 'set', 'uns', 'has']);
     }
 }
diff --git a/tests/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtensionUnitTest.php b/tests/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtensionUnitTest.php
new file mode 100644
index 0000000..1ce784c
--- /dev/null
+++ b/tests/bitExpert/PHPStan/Magento/Reflection/Framework/DataObjectMagicMethodReflectionExtensionUnitTest.php
@@ -0,0 +1,105 @@
+<?php
+
+/*
+ * This file is part of the phpstan-magento package.
+ *
+ * (c) bitExpert AG
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+declare(strict_types=1);
+
+namespace bitExpert\PHPStan\Magento\Reflection\Framework;
+
+use PHPStan\Reflection\ClassReflection;
+use PHPStan\Type\BooleanType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\StringType;
+use PHPStan\Type\UnionType;
+use PHPUnit\Framework\TestCase;
+
+class DataObjectMagicMethodReflectionExtensionUnitTest extends TestCase
+{
+    /**
+     * @var DataObjectMagicMethodReflectionExtension
+     */
+    private $extension;
+    /**
+     * @var ClassReflection|\PHPUnit\Framework\MockObject\MockObject
+     */
+    private $classReflection;
+
+    protected function setUp(): void
+    {
+        $this->extension = new DataObjectMagicMethodReflectionExtension();
+        $this->classReflection = $this->createMock(ClassReflection::class);
+    }
+
+    /**
+     * @test
+     */
+    public function returnMagicMethodReflectionForGetMethod(): void
+    {
+        $methodReflection = $this->extension->getMethod($this->classReflection, 'getTest');
+
+        $variants = $methodReflection->getVariants();
+        $params = $variants[0]->getParameters();
+
+        $this->assertCount(1, $variants);
+        $this->assertInstanceOf(MixedType::class, $variants[0]->getReturnType());
+        $this->assertCount(2, $params);
+        $this->assertInstanceOf(StringType::class, $params[0]->getType());
+        $this->assertInstanceOf(UnionType::class, $params[1]->getType());
+    }
+
+    /**
+     * @test
+     */
+    public function returnMagicMethodReflectionForSetMethod(): void
+    {
+        $methodReflection = $this->extension->getMethod($this->classReflection, 'setTest');
+
+        $variants = $methodReflection->getVariants();
+        $params = $variants[0]->getParameters();
+
+        $this->assertCount(1, $variants);
+        $this->assertInstanceOf(ObjectType::class, $variants[0]->getReturnType());
+        $this->assertCount(2, $params);
+        $this->assertInstanceOf(UnionType::class, $params[0]->getType());
+        $this->assertInstanceOf(MixedType::class, $params[1]->getType());
+    }
+
+    /**
+     * @test
+     */
+    public function returnMagicMethodReflectionForUnsetMethod(): void
+    {
+        $methodReflection = $this->extension->getMethod($this->classReflection, 'unsetTest');
+
+        $variants = $methodReflection->getVariants();
+        $params = $variants[0]->getParameters();
+
+        $this->assertCount(1, $variants);
+        $this->assertInstanceOf(ObjectType::class, $variants[0]->getReturnType());
+        $this->assertCount(1, $params);
+        $this->assertInstanceOf(UnionType::class, $params[0]->getType());
+    }
+
+    /**
+     * @test
+     */
+    public function returnMagicMethodReflectionForHasMethod(): void
+    {
+        $methodReflection = $this->extension->getMethod($this->classReflection, 'hasTest');
+
+        $variants = $methodReflection->getVariants();
+        $params = $variants[0]->getParameters();
+
+        $this->assertCount(1, $variants);
+        $this->assertInstanceOf(BooleanType::class, $variants[0]->getReturnType());
+        $this->assertCount(1, $params);
+        $this->assertInstanceOf(StringType::class, $params[0]->getType());
+    }
+}