Add test on adding raw stage

mongodb · Mar 18, 2024 · 537a828 · 537a828
1 parent 0abbaa9
commit 537a828
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 23 deletions.
diff --git a/src/Query/AggregationBuilder.php b/src/Query/AggregationBuilder.php
@@ -7,7 +7,8 @@
 use Illuminate\Support\Collection as LaravelCollection;
 use MongoDB\Builder\BuilderEncoder;
 use MongoDB\Builder\Stage\FluentFactoryTrait;
-use MongoDB\Laravel\Collection;
+use MongoDB\Collection as MongoDBCollection;
+use MongoDB\Laravel\Collection as LaravelMongoDBCollection;
 
 use function array_replace;
 use function collect;
@@ -17,9 +18,9 @@ final class AggregationBuilder
     use FluentFactoryTrait;
 
     public function __construct(
-        array $pipeline,
-        private Collection $collection,
-        private array $options,
+        private MongoDBCollection|LaravelMongoDBCollection $collection,
+        array $pipeline = [],
+        private array $options = [],
     ) {
         $this->pipeline = $pipeline;
     }

diff --git a/src/Query/Builder.php b/src/Query/Builder.php
@@ -290,7 +290,7 @@ private function getAggregationBuilder(): AggregationBuilder
             throw new BadMethodCallException('Aggregation builder requires package mongodb/builder 0.2+');
         }
 
-        $agg = new AggregationBuilder([], $this->collection, $this->options);
+        $agg = new AggregationBuilder($this->collection, [], $this->options);
 
         $wheres = $this->compileWheres();
 

diff --git a/tests/Query/AggregationBuilderTest.php b/tests/Query/AggregationBuilderTest.php
@@ -11,11 +11,15 @@
 use MongoDB\BSON\UTCDateTime;
 use MongoDB\Builder\BuilderEncoder;
 use MongoDB\Builder\Expression;
+use MongoDB\Builder\Pipeline;
 use MongoDB\Builder\Type\Sort;
+use MongoDB\Collection as MongoDBCollection;
 use MongoDB\Laravel\Query\AggregationBuilder;
 use MongoDB\Laravel\Tests\Models\User;
 use MongoDB\Laravel\Tests\TestCase;
 
+use function is_array;
+
 class AggregationBuilderTest extends TestCase
 {
     public function tearDown(): void
@@ -48,28 +52,20 @@ public function testCreateFromQueryBuilder(): void
             ->sort(year: Sort::Desc, name: Sort::Asc)
             ->unset('birthday');
 
-        // The encoder is used to convert the pipeline to a BSON document
-        $codec = new BuilderEncoder();
-        $json = Document::fromPHP([
-            'pipeline' => $codec->encode($pipeline->getPipeline()),
-        ])->toCanonicalExtendedJSON();
-
         // Compare with the expected pipeline
-        $expected = Document::fromPHP([
-            'pipeline' => [
-                ['$match' => ['name' => 'John Doe']],
-                ['$limit' => 10],
-                [
-                    '$addFields' => [
-                        'year' => ['$year' => ['date' => '$birthday']],
-                    ],
+        $expected = [
+            ['$match' => ['name' => 'John Doe']],
+            ['$limit' => 10],
+            [
+                '$addFields' => [
+                    'year' => ['$year' => ['date' => '$birthday']],
                 ],
-                ['$sort' => ['year' => -1, 'name' => 1]],
-                ['$unset' => ['birthday']],
             ],
-        ])->toCanonicalExtendedJSON();
+            ['$sort' => ['year' => -1, 'name' => 1]],
+            ['$unset' => ['birthday']],
+        ];
 
-        $this->assertJsonStringEqualsJsonString($expected, $json);
+        $this->assertSamePipeline($expected, $pipeline->getPipeline());
 
         // Execute the pipeline and validate the results
         $results = $pipeline->get();
@@ -81,4 +77,35 @@ public function testCreateFromQueryBuilder(): void
         $this->assertIsInt($results->first()['year']);
         $this->assertArrayNotHasKey('birthday', $results->first());
     }
+
+    public function testAddRawStage(): void
+    {
+        $collection = $this->createMock(MongoDBCollection::class);
+
+        $pipeline = new AggregationBuilder($collection);
+        $pipeline
+            ->addRawStage('$match', ['name' => 'John Doe'])
+            ->addRawStage('$limit', 10)
+            ->addRawStage('$replaceRoot', (object) ['newRoot' => '$$ROOT']);
+
+        $expected = [
+            ['$match' => ['name' => 'John Doe']],
+            ['$limit' => 10],
+            ['$replaceRoot' => ['$newRoot' => '$$ROOT']],
+        ];
+
+        $this->assertSamePipeline($expected, $pipeline->getPipeline());
+    }
+
+    private static function assertSamePipeline(array $expected, Pipeline $pipeline): void
+    {
+        $expected = Document::fromPHP(['pipeline' => $expected])->toCanonicalExtendedJSON();
+
+        $codec = new BuilderEncoder();
+        $actual = $codec->encode($pipeline);
+        // Normalize with BSON round-trip
+        $actual = Document::fromPHP(['pipeline' => $actual])->toCanonicalExtendedJSON();
+
+        self::assertJsonStringEqualsJsonString($expected, $actual);
+    }
 }