Add assertion and refactor

Author: GromNaN

generator/src/FluentStageFactoryGenerator.php

@@ -30,9 +30,7 @@
 use Nette\PhpGenerator\Method;
 use Nette\PhpGenerator\PhpNamespace;
 use Nette\PhpGenerator\TraitType;
-use RuntimeException;
 use stdClass;
-use Throwable;
 
 use function array_key_last;
 use function array_map;
@@ -99,36 +97,26 @@ private function createFluentFactoryClass(GeneratorDefinition $definition): PhpN
         $staticFactory = ClassType::from(self::FACTORY_CLASS);
         assert($staticFactory instanceof ClassType);
         foreach ($staticFactory->getMethods() as $method) {
-            if (! $method->isPublic()) {
-                continue;
-            }
-
-            try {
-                $this->addMethod($method, $namespace, $class);
-            } catch (Throwable $e) {
-                throw new RuntimeException(sprintf('Failed to generate class for operator "%s"', $operator->name), 0, $e);
-            }
+            $this->addMethod($method, $class);
         }
 
-        $staticFactory = TraitType::from(Stage\FactoryTrait::class);
-        assert($staticFactory instanceof TraitType);
-        foreach ($staticFactory->getMethods() as $method) {
-            if (! $method->isPublic()) {
-                continue;
-            }
-
-            try {
-                $this->addMethod($method, $namespace, $class);
-            } catch (Throwable $e) {
-                throw new RuntimeException(sprintf('Failed to generate class for operator "%s"', $operator->name), 0, $e);
+        foreach ($staticFactory->getTraits() as $trait) {
+            $staticFactory = TraitType::from($trait->getName());
+            assert($staticFactory instanceof TraitType);
+            foreach ($staticFactory->getMethods() as $method) {
+                $this->addMethod($method, $class);
             }
         }
 
         return $namespace;
     }
 
-    private function addMethod(Method $factoryMethod, PhpNamespace $namespace, ClassType $class): void
+    private function addMethod(Method $factoryMethod, ClassType $class): void
     {
+        if (! $factoryMethod->isPublic()) {
+            return;
+        }
+
         if ($class->hasMethod($factoryMethod->getName())) {
             return;
         }

tests/Builder/FluentPipelineFactoryTest.php

@@ -8,9 +8,8 @@
 use MongoDB\Builder\Query;
 use MongoDB\Builder\Stage\FluentFactory;
 use MongoDB\Builder\Type\Sort;
-use PHPUnit\Framework\TestCase;
 
-class FluentPipelineFactoryTest extends TestCase
+class FluentPipelineFactoryTest extends PipelineTestCase
 {
     public function testFluentPipelineFactory(): void
     {
@@ -22,5 +21,13 @@ public function testFluentPipelineFactory(): void
 
         $this->assertInstanceof(Pipeline::class, $pipeline);
         $this->assertCount(3, $pipeline->getIterator());
+
+        $this->assertSamePipeline(<<<'JSON'
+        [
+            { "$match": {"x": {"$eq": {"$numberInt": "1"}}}},
+            { "$project": {"_id": false, "x": true}},
+            { "$sort": {"x": {"$numberInt": "1"}}}
+        ]
+        JSON, $pipeline);
     }
 }