diff --git a/tests/Doctrine/Tests/Models/ECommerce/ECommerceCategory.php b/tests/Doctrine/Tests/Models/ECommerce/ECommerceCategory.php
index b7c7b4a68..af4e088d5 100644
--- a/tests/Doctrine/Tests/Models/ECommerce/ECommerceCategory.php
+++ b/tests/Doctrine/Tests/Models/ECommerce/ECommerceCategory.php
@@ -59,9 +59,9 @@ class ECommerceCategory
 
     public function removeProduct(ECommerceProduct $product)
     {
-        if ($this->products->contains($product)) {
-            $this->products->removeElement($product);
-            $product->removeCategory($this);
+        $removed = $this->products->removeElement($product);
+        if ($removed !== null) {
+            $removed->removeCategory($this);
         }
     }
 
diff --git a/tests/Doctrine/Tests/Models/ECommerce/ECommerceProduct.php b/tests/Doctrine/Tests/Models/ECommerce/ECommerceProduct.php
index 2a457ce52..3bc4b0b20 100644
--- a/tests/Doctrine/Tests/Models/ECommerce/ECommerceProduct.php
+++ b/tests/Doctrine/Tests/Models/ECommerce/ECommerceProduct.php
@@ -98,12 +98,10 @@ class ECommerceProduct
 
     public function removeFeature(ECommerceFeature $feature)
     {
-        if ($this->features->contains($feature)) {
-            $removed = $this->features->removeElement($feature);
-            if ($removed) {
-                $feature->removeProduct();
-                return true;
-            }
+        $removed = $this->features->removeElement($feature);
+        if ($removed !== null) {
+            $removed->removeProduct();
+            return true;
         }
         return false;
     }
@@ -118,9 +116,9 @@ class ECommerceProduct
 
     public function removeCategory(ECommerceCategory $category)
     {
-        if ($this->categories->contains($category)) {
-            $this->categories->removeElement($category);
-            $category->removeProduct($this);
+        $removed = $this->categories->removeElement($category);
+        if ($removed !== null) {
+            $removed->removeProduct($this);
         }
     }