From b52c3735ba97ac7c7dca339ec7b9501ff3236ff9 Mon Sep 17 00:00:00 2001
From: phleauran <f7102v+github@gmail.com>
Date: Fri, 1 Nov 2024 20:20:53 +0000
Subject: [PATCH] Add profiling for countDocuments() and
 estimatedDocumentCount()

---
 src/Capsule/Collection.php                  | 24 +++++++++++++++++++++
 tests/Functional/Capsule/CollectionTest.php | 22 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/src/Capsule/Collection.php b/src/Capsule/Collection.php
index ad51e13..fdf3074 100644
--- a/src/Capsule/Collection.php
+++ b/src/Capsule/Collection.php
@@ -60,6 +60,18 @@ public function count($filter = [], array $options = [])
         return $result;
     }
 
+    /**
+     * @inheritDoc
+     */
+    public function countDocuments($filter = [], array $options = [])
+    {
+        $query = $this->prepareQuery(__FUNCTION__, $filter, [], $options);
+        $result = parent::countDocuments($query->getFilters(), $query->getOptions());
+        $this->notifyQueryExecution($query);
+
+        return $result;
+    }
+
     /**
      * @inheritDoc
      */
@@ -180,6 +192,18 @@ public function distinct($fieldName, $filter = [], array $options = [])
         return $result;
     }
 
+    /**
+     * @inheritDoc
+     */
+    public function estimatedDocumentCount(array $options = [])
+    {
+        $query = $this->prepareQuery(__FUNCTION__, [], [], $options);
+        $result = parent::estimatedDocumentCount($options);
+        $this->notifyQueryExecution($query);
+
+        return $result;
+    }
+
     /**
      * @param array|object $filters
      * @param array|object $data
diff --git a/tests/Functional/Capsule/CollectionTest.php b/tests/Functional/Capsule/CollectionTest.php
index 69f4bf9..cf4c3cc 100644
--- a/tests/Functional/Capsule/CollectionTest.php
+++ b/tests/Functional/Capsule/CollectionTest.php
@@ -71,6 +71,17 @@ public function test_count(): void
         $coll->count(['test' => 1]);
     }
 
+    public function test_countDocuments(): void
+    {
+        $manager = $this->getManager();
+        $ev = $this->prophesize(EventDispatcherInterface::class);
+        $this->assertEventsDispatching($ev);
+
+        $coll = new Collection($manager, 'test_client', 'testdb', 'test_collection', [], $ev->reveal());
+
+        $coll->countDocuments(['test' => 1]);
+    }
+
     public function test_find(): void
     {
         $manager = $this->getManager();
@@ -190,6 +201,17 @@ public function test_distinct(): void
         $coll->distinct('field');
     }
 
+    public function test_estimatedDocumentCount(): void
+    {
+        $manager = $this->getManager();
+        $ev = $this->prophesize(EventDispatcherInterface::class);
+        $this->assertEventsDispatching($ev);
+
+        $coll = new Collection($manager, 'test_client', 'testdb', 'test_document', [], $ev->reveal());
+
+        $coll->estimatedDocumentCount();
+    }
+
     protected function assertEventsDispatching($ev)
     {
         $ev->dispatch(Argument::type(QueryEvent::class), QueryEvent::QUERY_PREPARED)->shouldBeCalled();