Fix AssertSameWithCountRule for recursive count()

Author: staabm

src/Rules/PHPUnit/AssertSameWithCountRule.php

@@ -8,8 +8,12 @@
 use PHPStan\Analyser\Scope;
 use PHPStan\Rules\Rule;
 use PHPStan\Rules\RuleErrorBuilder;
+use PHPStan\TrinaryLogic;
+use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
 use function count;
+use const COUNT_NORMAL;
 
 /**
  * @implements Rule<CallLike>
@@ -42,36 +46,67 @@ public function processNode(Node $node, Scope $scope): array
 		}
 
 		$right = $node->getArgs()[1]->value;
-		if (
-			$right instanceof Node\Expr\FuncCall
-			&& $right->name instanceof Node\Name
-			&& $right->name->toLowerString() === 'count'
-		) {
+		if (self::isCountFunctionCall($right, $scope)) {
 			return [
 				RuleErrorBuilder::message('You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, count($variable)).')
 					->identifier('phpunit.assertCount')
 					->build(),
 			];
 		}
 
+		if (self::isCountableMethodCall($right, $scope)) {
+			return [
+				RuleErrorBuilder::message('You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, $variable->count()).')
+					->identifier('phpunit.assertCount')
+					->build(),
+			];
+		}
+
+		return [];
+	}
+
+	/**
+	 * @phpstan-assert-if-true Node\Expr\FuncCall $expr
+	 */
+	private static function isCountFunctionCall(Node\Expr $expr, Scope $scope): bool
+	{
+		return $expr instanceof Node\Expr\FuncCall
+			&& $expr->name instanceof Node\Name
+			&& $expr->name->toLowerString() === 'count'
+			&& count($expr->getArgs()) >= 1
+			&& self::isNormalCount($expr, $scope->getType($expr->getArgs()[0]->value), $scope)->yes();
+	}
+
+	/**
+	 * @phpstan-assert-if-true Node\Expr\MethodCall $expr
+	 */
+	private static function isCountableMethodCall(Node\Expr $expr, Scope $scope): bool
+	{
 		if (
-			$right instanceof Node\Expr\MethodCall
-			&& $right->name instanceof Node\Identifier
-			&& $right->name->toLowerString() === 'count'
-			&& count($right->getArgs()) === 0
+			$expr instanceof Node\Expr\MethodCall
+			&& $expr->name instanceof Node\Identifier
+			&& $expr->name->toLowerString() === 'count'
+			&& count($expr->getArgs()) === 0
 		) {
-			$type = $scope->getType($right->var);
+			$type = $scope->getType($expr->var);
 
 			if ((new ObjectType(Countable::class))->isSuperTypeOf($type)->yes()) {
-				return [
-					RuleErrorBuilder::message('You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, $variable->count()).')
-						->identifier('phpunit.assertCount')
-						->build(),
-				];
+				return true;
 			}
 		}
 
-		return [];
+		return false;
+	}
+
+	private static function isNormalCount(Node\Expr\FuncCall $countFuncCall, Type $countedType, Scope $scope): TrinaryLogic
+	{
+		if (count($countFuncCall->getArgs()) === 1) {
+			$isNormalCount = TrinaryLogic::createYes();
+		} else {
+			$mode = $scope->getType($countFuncCall->getArgs()[1]->value);
+			$isNormalCount = (new ConstantIntegerType(COUNT_NORMAL))->isSuperTypeOf($mode)->result->or($countedType->getIterableValueType()->isArray()->negate());
+		}
+		return $isNormalCount;
 	}
 
 }

tests/Rules/PHPUnit/AssertSameWithCountRuleTest.php

@@ -31,6 +31,14 @@ public function testRule(): void
 				'You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, $variable->count()).',
 				30,
 			],
+			[
+				'You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, count($variable)).',
+				40,
+			],
+			[
+				'You should use assertCount($expectedCount, $variable) instead of assertSame($expectedCount, count($variable)).',
+				45,
+			],
 		]);
 	}
 

tests/Rules/PHPUnit/data/assert-same-count.php

@@ -30,11 +30,35 @@ public function testAssertSameWithCountMethodForCountableVariableIsNotOK()
 		$this->assertSame(5, $foo->bar->count());
 	}
 
+	public function testRecursiveCount($x)
+	{
+		$this->assertSame(5, count([1, 2, 3, $x], COUNT_RECURSIVE)); // OK
+	}
+
+	public function testNormalCount($x)
+	{
+		$this->assertSame(5, count([1, 2, 3, $x], COUNT_NORMAL));
+	}
+
+	public function testImplicitNormalCount($mode)
+	{
+		$this->assertSame(5, count([1, 2, 3], $mode));
+	}
+
+	public function testUnknownCountable($x, $mode)
+	{
+		$this->assertSame(5, count($x, $mode)); // OK
+	}
+
+	public function testUnknownCountMode($x, $mode)
+	{
+		$this->assertSame(5, count([1, 2, 3, $x], $mode)); // OK
+	}
 }
 
 class Bar implements \Countable {
 	public function count(): int
 	{
 		return 1;
 	}
-};
+}