Set correct types for mocks of multiple classes

extension.neon

@@ -5,9 +5,7 @@ parameters:
 			- markTestIncomplete
 			- markTestSkipped
 	stubFiles:
-		- stubs/MockBuilder.stub
 		- stubs/MockObject.stub
-		- stubs/TestCase.stub
 
 services:
 	-
@@ -26,6 +24,14 @@ services:
 		class: PHPStan\Type\PHPUnit\Assert\AssertStaticMethodTypeSpecifyingExtension
 		tags:
 			- phpstan.typeSpecifier.staticMethodTypeSpecifyingExtension
+	-
+		class: PHPStan\Type\PHPUnit\CreateMockDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: PHPStan\Type\PHPUnit\GetMockBuilderDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
 		class: PHPStan\Type\PHPUnit\MockBuilderDynamicReturnTypeExtension
 		tags:

phpstan.neon

@@ -7,6 +7,7 @@ includes:
 parameters:
 	excludes_analyse:
 		- tests/*/data/*
+	treatPhpDocTypesAsCertain: false
 
 services:
 	scopeIsInClass:

src/Type/PHPUnit/CreateMockDynamicReturnTypeExtension.php

@@ -0,0 +1,70 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\PHPUnit;
+
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+
+class CreateMockDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	/** @var int[] */
+	private $methods = [
+		'createMock' => 0,
+		'createConfiguredMock' => 0,
+		'createPartialMock' => 0,
+		'createTestProxy' => 0,
+		'getMockForAbstractClass' => 0,
+		'getMockFromWsdl' => 1,
+	];
+
+	public function getClass(): string
+	{
+		return 'PHPUnit\Framework\TestCase';
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		$name = $methodReflection->getName();
+		return array_key_exists($methodReflection->getName(), $this->methods);
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		$argumentIndex = $this->methods[$methodReflection->getName()];
+		$parametersAcceptor = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants());
+		if (!isset($methodCall->args[$argumentIndex])) {
+			return $parametersAcceptor->getReturnType();
+		}
+		$argType = $scope->getType($methodCall->args[$argumentIndex]->value);
+
+		$types = [];
+		if ($argType instanceof ConstantStringType) {
+			$types[] = new ObjectType($argType->getValue());
+		}
+
+		if ($argType instanceof ConstantArrayType) {
+			$types = array_map(function (ConstantStringType $argType): ObjectType {
+				return new ObjectType($argType->getValue());
+			}, $argType->getValueTypes());
+		}
+
+		if (count($types) === 0) {
+			return $parametersAcceptor->getReturnType();
+		}
+
+		return TypeCombinator::intersect(
+			$parametersAcceptor->getReturnType(),
+			...$types
+		);
+	}
+
+}

src/Type/PHPUnit/GetMockBuilderDynamicReturnTypeExtension.php

@@ -0,0 +1,57 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\PHPUnit;
+
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeWithClassName;
+
+class GetMockBuilderDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	public function getClass(): string
+	{
+		return 'PHPUnit\Framework\TestCase';
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'getMockBuilder';
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		$parametersAcceptor = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants());
+		$mockBuilderType = $parametersAcceptor->getReturnType();
+		if (count($methodCall->args) === 0) {
+			return $mockBuilderType;
+		}
+		if (!$mockBuilderType instanceof TypeWithClassName) {
+			throw new \PHPStan\ShouldNotHappenException();
+		}
+
+		$argType = $scope->getType($methodCall->args[0]->value);
+		if ($argType instanceof ConstantStringType) {
+			$class = $argType->getValue();
+
+			return new MockBuilderType($mockBuilderType, $class);
+		}
+
+		if ($argType instanceof ConstantArrayType) {
+			$classes = array_map(function (ConstantStringType $argType): string {
+				return $argType->getValue();
+			}, $argType->getValueTypes());
+
+			return new MockBuilderType($mockBuilderType, ...$classes);
+		}
+
+		return $mockBuilderType;
+	}
+
+}

src/Type/PHPUnit/MockBuilderDynamicReturnTypeExtension.php

@@ -4,34 +4,72 @@
 
 use PhpParser\Node\Expr\MethodCall;
 use PHPStan\Analyser\Scope;
+use PHPStan\Broker\Broker;
+use PHPStan\Reflection\BrokerAwareExtension;
 use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
-use PHPUnit\Framework\MockObject\MockBuilder;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeWithClassName;
 
-class MockBuilderDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+class MockBuilderDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension, BrokerAwareExtension
 {
 
+	/** @var \PHPStan\Broker\Broker */
+	private $broker;
+
+	public function setBroker(Broker $broker): void
+	{
+		$this->broker = $broker;
+	}
+
 	public function getClass(): string
 	{
-		return MockBuilder::class;
+		$testCase = $this->broker->getClass('PHPUnit\Framework\TestCase');
+		$mockBuilderType = ParametersAcceptorSelector::selectSingle(
+			$testCase->getNativeMethod('getMockBuilder')->getVariants()
+		)->getReturnType();
+		if (!$mockBuilderType instanceof TypeWithClassName) {
+			throw new \PHPStan\ShouldNotHappenException();
+		}
+
+		return $mockBuilderType->getClassName();
 	}
 
 	public function isMethodSupported(MethodReflection $methodReflection): bool
 	{
-		return !in_array(
+		return true;
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		$calledOnType = $scope->getType($methodCall->var);
+		if (!in_array(
 			$methodReflection->getName(),
 			[
 				'getMock',
 				'getMockForAbstractClass',
-				'getMockForTrait',
 			],
 			true
-		);
-	}
+		)) {
+			return $calledOnType;
+		}
 
-	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
-	{
-		return $scope->getType($methodCall->var);
+		$parametersAcceptor = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants());
+
+		if (!$calledOnType instanceof MockBuilderType) {
+			return $parametersAcceptor->getReturnType();
+		}
+		$types = array_map(function (string $type): ObjectType {
+			return new ObjectType($type);
+		}, $calledOnType->getMockedClasses());
+
+		return TypeCombinator::intersect(
+			$parametersAcceptor->getReturnType(),
+			...$types
+		);
 	}
 
 }

src/Type/PHPUnit/MockBuilderType.php

@@ -0,0 +1,37 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\PHPUnit;
+
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\TypeWithClassName;
+use PHPStan\Type\VerbosityLevel;
+
+class MockBuilderType extends ObjectType
+{
+
+	/** @var array<string> */
+	private $mockedClasses;
+
+	public function __construct(
+		TypeWithClassName $mockBuilderType,
+		string ...$mockedClasses
+	)
+	{
+		parent::__construct($mockBuilderType->getClassName());
+		$this->mockedClasses = $mockedClasses;
+	}
+
+	/**
+	 * @return array<string>
+	 */
+	public function getMockedClasses(): array
+	{
+		return $this->mockedClasses;
+	}
+
+	public function describe(VerbosityLevel $level): string
+	{
+		return sprintf('%s<%s>', parent::describe($level), implode('&', $this->mockedClasses));
+	}
+
+}