diff --git a/stubs/Collections/Collection.stub b/stubs/Collections/Collection.stub
index dbf36bd6..f6aafc8c 100644
--- a/stubs/Collections/Collection.stub
+++ b/stubs/Collections/Collection.stub
@@ -16,4 +16,31 @@ use IteratorAggregate;
 interface Collection extends Countable, IteratorAggregate, ArrayAccess, ReadableCollection
 {
 
+    /**
+     * @phpstan-impure
+     *
+     * @param T $element
+     *
+     * @return true
+     */
+    public function add($element) {}
+
+    /**
+     * @phpstan-impure
+     *
+     * @param TKey $key
+     *
+     * @return T|null
+     */
+    public function remove($key) {}
+
+    /**
+     * @phpstan-impure
+     *
+     * @param T $element
+     *
+     * @return bool
+     */
+    public function removeElement($element) {}
+
 }
diff --git a/tests/Type/Doctrine/Collection/IsEmptyTypeSpecifyingExtensionTest.php b/tests/Type/Doctrine/Collection/IsEmptyTypeSpecifyingExtensionTest.php
index 1678253b..83fc20c9 100644
--- a/tests/Type/Doctrine/Collection/IsEmptyTypeSpecifyingExtensionTest.php
+++ b/tests/Type/Doctrine/Collection/IsEmptyTypeSpecifyingExtensionTest.php
@@ -2,48 +2,28 @@
 
 namespace PHPStan\Type\Doctrine\Collection;
 
-use PHPStan\Rules\Rule;
-use PHPStan\Testing\RuleTestCase;
+use PHPStan\Testing\TypeInferenceTestCase;
 
-/**
- * @extends RuleTestCase<VariableTypeReportingRule>
- */
-class IsEmptyTypeSpecifyingExtensionTest extends RuleTestCase
+class IsEmptyTypeSpecifyingExtensionTest extends TypeInferenceTestCase
 {
 
-	protected function getRule(): Rule
+	/** @return iterable<mixed> */
+	public function dataFileAsserts(): iterable
 	{
-		return new VariableTypeReportingRule();
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/collection.php');
 	}
 
-	public function testExtension(): void
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param mixed ...$args
+	 */
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args
+	): void
 	{
-		$this->analyse([__DIR__ . '/data/collection.php'], [
-			[
-				'Variable $entityOrFalse1 is: MyEntity|false',
-				18,
-			],
-			[
-				'Variable $entityOrFalse2 is: MyEntity|false',
-				21,
-			],
-			[
-				'Variable $false1 is: false',
-				25,
-			],
-			[
-				'Variable $false2 is: false',
-				28,
-			],
-			[
-				'Variable $entity1 is: MyEntity',
-				33,
-			],
-			[
-				'Variable $entity2 is: MyEntity',
-				36,
-			],
-		]);
+		$this->assertFileAsserts($assertType, $file, ...$args);
 	}
 
 	public static function getAdditionalConfigFiles(): array
diff --git a/tests/Type/Doctrine/Collection/data/collection.php b/tests/Type/Doctrine/Collection/data/collection.php
index 249a8152..14a6c95b 100644
--- a/tests/Type/Doctrine/Collection/data/collection.php
+++ b/tests/Type/Doctrine/Collection/data/collection.php
@@ -2,6 +2,8 @@
 
 use Doctrine\Common\Collections\ArrayCollection;
 
+use function PHPStan\Testing\assertType;
+
 class MyEntity
 {
 
@@ -15,23 +17,41 @@ class MyEntity
 $collection = new ArrayCollection();
 
 $entityOrFalse1 = $collection->first();
-$entityOrFalse1;
+assertType('MyEntity|false', $entityOrFalse1);
 
 $entityOrFalse2 = $collection->last();
-$entityOrFalse2;
+assertType('MyEntity|false', $entityOrFalse2);
 
 if ($collection->isEmpty()) {
 	$false1 = $collection->first();
-	$false1;
+	assertType('false', $false1);
 
 	$false2 = $collection->last();
-	$false2;
+	assertType('false', $false2);
 }
 
 if (!$collection->isEmpty()) {
 	$entity1 = $collection->first();
-	$entity1;
+	assertType('MyEntity', $entity1);
 
 	$entity2 = $collection->last();
-	$entity2;
+	assertType('MyEntity', $entity2);
+}
+
+if ($collection->isEmpty()) {
+	$collection->add($new);
+	$result1 = $collection->first();
+	assertType('MyEntity|false', $result1);
+
+	$result2 = $collection->last();
+	assertType('MyEntity|false', $result2);
+}
+
+if (!$collection->isEmpty()) {
+	$collection->removeElement($new);
+	$result3 = $collection->first();
+	assertType('MyEntity|false', $result3);
+
+	$result4 = $collection->last();
+	assertType('MyEntity|false', $result4);
 }