diff --git a/src/Type/PHPUnit/MockBuilderDynamicReturnTypeExtension.php b/src/Type/PHPUnit/MockBuilderDynamicReturnTypeExtension.php
index b75c69f..d85af32 100644
--- a/src/Type/PHPUnit/MockBuilderDynamicReturnTypeExtension.php
+++ b/src/Type/PHPUnit/MockBuilderDynamicReturnTypeExtension.php
@@ -49,6 +49,7 @@ public function getTypeFromMethodCall(MethodReflection $methodReflection, Method
 			[
 				'getMock',
 				'getMockForAbstractClass',
+				'getMockForTrait',
 			],
 			true
 		)) {
diff --git a/tests/Rules/PHPUnit/AssertSameMethodDifferentTypesRuleTest.php b/tests/Rules/PHPUnit/AssertSameMethodDifferentTypesRuleTest.php
index dc3bb7e..95128c4 100644
--- a/tests/Rules/PHPUnit/AssertSameMethodDifferentTypesRuleTest.php
+++ b/tests/Rules/PHPUnit/AssertSameMethodDifferentTypesRuleTest.php
@@ -12,7 +12,7 @@ class AssertSameMethodDifferentTypesRuleTest extends \PHPStan\Testing\RuleTestCa
 
 	protected function getRule(): Rule
 	{
-		return new ImpossibleCheckTypeMethodCallRule(new ImpossibleCheckTypeHelper($this->getTypeSpecifier()), true);
+		return new ImpossibleCheckTypeMethodCallRule(new ImpossibleCheckTypeHelper($this->createBroker(), $this->getTypeSpecifier()), true);
 	}
 
 	/**
diff --git a/tests/Rules/PHPUnit/AssertSameStaticMethodDifferentTypesRuleTest.php b/tests/Rules/PHPUnit/AssertSameStaticMethodDifferentTypesRuleTest.php
index 7fbda4f..977b62c 100644
--- a/tests/Rules/PHPUnit/AssertSameStaticMethodDifferentTypesRuleTest.php
+++ b/tests/Rules/PHPUnit/AssertSameStaticMethodDifferentTypesRuleTest.php
@@ -12,7 +12,7 @@ class AssertSameStaticMethodDifferentTypesRuleTest extends \PHPStan\Testing\Rule
 
 	protected function getRule(): Rule
 	{
-		return new ImpossibleCheckTypeStaticMethodCallRule(new ImpossibleCheckTypeHelper($this->getTypeSpecifier()), true);
+		return new ImpossibleCheckTypeStaticMethodCallRule(new ImpossibleCheckTypeHelper($this->createBroker(), $this->getTypeSpecifier()), true);
 	}
 
 	/**