Provides correct class for method

README.md

@@ -25,6 +25,8 @@
 * Provides correct attributes for `ObjectBehavior`:
   * public attributes
   * static properties (with `$this->CONSTANT_NAME`)
+* Provides correct class for `getWrappedObject` method
+
 
 ## Compatibility
 

spec/Proget/Tests/FooSpec.php

@@ -21,6 +21,10 @@ public function it_should_return_int_from_baz(Bar $bar, Baz $baz): void
         $bar->baz = $baz;
         $baz->someInt()->willReturn(99);
 
-        $this->getIntFromBaz($bar)->shouldBe(99);
+        // intentionally to test support for correct return type
+        $int = $baz->getWrappedObject()->someInt();
+        $this->getWrappedObject()->property;
+
+        $this->getIntFromBaz($bar)->shouldBe($int);
     }
 }

src/Reflection/GetWrappedObjectMethodReflection.php

@@ -0,0 +1,73 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Proget\PHPStan\PhpSpec\Reflection;
+
+use PHPStan\Reflection\ClassMemberReflection;
+use PHPStan\Reflection\ClassReflection;
+use PHPStan\Reflection\FunctionVariant;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\ObjectType;
+
+final class GetWrappedObjectMethodReflection implements MethodReflection
+{
+    /**
+     * @var MethodReflection
+     */
+    private $methodReflection;
+
+    /**
+     * @var string
+     */
+    private $wrappedClass;
+
+    public function __construct(MethodReflection $methodReflection, string $wrappedClass)
+    {
+        $this->methodReflection = $methodReflection;
+        $this->wrappedClass = $wrappedClass;
+    }
+
+    public function getDeclaringClass(): ClassReflection
+    {
+        return $this->methodReflection->getDeclaringClass();
+    }
+
+    public function isStatic(): bool
+    {
+        return $this->methodReflection->isStatic();
+    }
+
+    public function isPrivate(): bool
+    {
+        return $this->methodReflection->isPrivate();
+    }
+
+    public function isPublic(): bool
+    {
+        return $this->methodReflection->isPublic();
+    }
+
+    public function getName(): string
+    {
+        return $this->methodReflection->getName();
+    }
+
+    public function getPrototype(): ClassMemberReflection
+    {
+        return $this->methodReflection->getPrototype();
+    }
+
+    public function getVariants(): array
+    {
+        $variant = $this->methodReflection->getVariants()[0];
+
+        return [
+            new FunctionVariant(
+                $variant->getParameters(),
+                $variant->isVariadic(),
+                new ObjectType($this->wrappedClass)
+            )
+        ];
+    }
+}

src/Reflection/ObjectBehaviorMethodsClassReflectionExtension.php

@@ -57,10 +57,14 @@ public function getMethod(ClassReflection $classReflection, string $methodName):
             return $subjectReflection->getMethod($methodName, new OutOfClassScope());
         }
 
-        return $this->getWrappedClassMethodReflection($classReflection, $methodName);
+        return new ObjectBehaviorMethodReflection(
+            $this->broker->getClass(
+                $this->getSourceClassName($classReflection)
+            )->getMethod($methodName, new OutOfClassScope())
+        );
     }
 
-    private function getWrappedClassMethodReflection(ClassReflection $classReflection, $methodName): MethodReflection
+    private function getSourceClassName(ClassReflection $classReflection): string
     {
         /** @var PSR0Resource[] $resources */
         $resources = $this->locator->findResources((string) $classReflection->getFileName());
@@ -74,8 +78,6 @@ private function getWrappedClassMethodReflection(ClassReflection $classReflectio
             throw new SpecSourceClassNotFound(sprintf('Spec source class %s not found', $className));
         }
 
-        $srcClassReflection = $this->broker->getClass($className);
-
-        return new ObjectBehaviorMethodReflection($srcClassReflection->getMethod($methodName, new OutOfClassScope()));
+        return $className;
     }
 }

src/Reflection/SpoofedCollaboratorMethodsClassReflectionExtension.php

@@ -35,13 +35,17 @@ public function hasMethod(ClassReflection $classReflection, string $methodName):
 
     public function getMethod(ClassReflection $classReflection, string $methodName): MethodReflection
     {
+        $collaboratorClassName = (string) preg_replace('/Collaborator$/', '', SpoofedCollaboratorRegistry::getAlias($classReflection->getName()));
         $collaboratorReflection = $this->broker->getClass(Collaborator::class);
+
+        if ($methodName === 'getWrappedObject') {
+            return new GetWrappedObjectMethodReflection($collaboratorReflection->getMethod($methodName, new OutOfClassScope()), $collaboratorClassName);
+        }
+
         if ($collaboratorReflection->hasMethod($methodName)) {
             return $collaboratorReflection->getMethod($methodName, new OutOfClassScope());
         }
 
-        $collaboratorClassName = (string) preg_replace('/Collaborator$/', '', SpoofedCollaboratorRegistry::getAlias($classReflection->getName()));
-
         return new CollaboratorMethodReflection($this->broker->getClass($collaboratorClassName)->getMethod($methodName, new OutOfClassScope()));
     }
 }

src/Type/ObjectBehaviorDynamicMethodReturnTypeExtension.php

@@ -6,32 +6,51 @@
 
 use PhpParser\Node\Expr\MethodCall;
 use PhpSpec\Locator\PSR0\PSR0Locator;
+use PhpSpec\Locator\PSR0\PSR0Resource;
 use PhpSpec\Locator\Resource;
+use PhpSpec\Locator\ResourceLocator;
 use PhpSpec\ObjectBehavior;
+use PhpSpec\Util\Filesystem;
 use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\ClassReflection;
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\ResourceType;
 use PHPStan\Type\Type;
+use Proget\PHPStan\PhpSpec\Exception\SpecSourceClassNotFound;
 use Proget\PHPStan\PhpSpec\Reflection\ObjectBehaviorMethodReflection;
 
 final class ObjectBehaviorDynamicMethodReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
+    /**
+     * @var ResourceLocator
+     */
+    private $locator;
+
+    public function __construct()
+    {
+        $this->locator = new PSR0Locator(new Filesystem());
+    }
+
     public function getClass(): string
     {
         return ObjectBehavior::class;
     }
 
     public function isMethodSupported(MethodReflection $methodReflection): bool
     {
-        return $methodReflection instanceof ObjectBehaviorMethodReflection;
+        return $methodReflection instanceof ObjectBehaviorMethodReflection || $methodReflection->getName() === 'getWrappedObject';
     }
 
     public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
     {
+        if ($methodReflection->getName() === 'getWrappedObject') {
+            return new ObjectType($this->getSourceClassName($scope->getClassReflection()));
+        }
+
         if (!$methodReflection instanceof ObjectBehaviorMethodReflection) {
             throw new ShouldNotHappenException();
         }
@@ -50,4 +69,26 @@ public function getTypeFromMethodCall(MethodReflection $methodReflection, Method
 
         return new SubjectType($returnType);
     }
+
+    // todo: looks like duplication from ObjectBehaviorMethodsClassReflectionExtension
+    private function getSourceClassName(?ClassReflection $classReflection): string
+    {
+        if ($classReflection === null) {
+            throw new ShouldNotHappenException();
+        }
+
+        /** @var PSR0Resource[] $resources */
+        $resources = $this->locator->findResources((string) $classReflection->getFileName());
+
+        if (count($resources) === 0) {
+            throw new SpecSourceClassNotFound(sprintf('Source class from %s not found', $classReflection->getName()));
+        }
+
+        $className = $resources[0]->getSrcClassname();
+        if (!class_exists($className)) {
+            throw new SpecSourceClassNotFound(sprintf('Spec source class %s not found', $className));
+        }
+
+        return $className;
+    }
 }