Extract fail point management to a trait used by Context

Author: jmikola

tests/UnifiedSpecTests/Context.php

@@ -5,7 +5,6 @@
 use LogicException;
 use MongoDB\Client;
 use MongoDB\Driver\ClientEncryption;
-use MongoDB\Driver\Server;
 use MongoDB\Driver\ServerApi;
 use MongoDB\Model\BSONArray;
 use MongoDB\Tests\FunctionalTestCase;
@@ -39,6 +38,8 @@
  */
 final class Context
 {
+    use ManagesFailPointsTrait;
+
     private ?string $activeClient = null;
 
     private EntityMap $entityMap;
@@ -61,9 +62,6 @@ final class Context
 
     private ?object $advanceClusterTime = null;
 
-    /** @var list<array{failPoint: stdClass, server: Server}> */
-    private array $failPoints = [];
-
     public function __construct(Client $internalClient, string $uri)
     {
         $this->entityMap = new EntityMap();
@@ -236,20 +234,6 @@ public function stopEventCollectors(): void
         }
     }
 
-    public function registerFailPoint(stdClass $failPoint, Server $server): void
-    {
-        $this->failPoints[] = [
-            'failPoint' => $failPoint,
-            'server' => $server,
-        ];
-    }
-
-    /** @return list<array{failPoint: stdClass, server: Server}> */
-    public function getFailPoints(): array
-    {
-        return $this->failPoints;
-    }
-
     /** @param string|array $readPreferenceTags */
     private function convertReadPreferenceTags($readPreferenceTags): array
     {

tests/UnifiedSpecTests/ManagesFailPointsTrait.php

@@ -0,0 +1,40 @@
+<?php
+
+namespace MongoDB\Tests\UnifiedSpecTests;
+
+use MongoDB\Driver\Server;
+use MongoDB\Operation\DatabaseCommand;
+use stdClass;
+
+use function PHPUnit\Framework\assertIsString;
+use function PHPUnit\Framework\assertObjectHasAttribute;
+
+trait ManagesFailPointsTrait
+{
+    /** @var list<list{string, Server}> */
+    private array $failPointsAndServers = [];
+
+    public function configureFailPoint(stdClass $failPoint, Server $server): void
+    {
+        assertObjectHasAttribute('configureFailPoint', $failPoint);
+        assertIsString($failPoint->configureFailPoint);
+        assertObjectHasAttribute('mode', $failPoint);
+
+        $operation = new DatabaseCommand('admin', $failPoint);
+        $operation->execute($server);
+
+        if ($failPoint->mode !== 'off') {
+            $this->failPointsAndServers[] = [$failPoint->configureFailPoint, $server];
+        }
+    }
+
+    public function disableFailPoints(): void
+    {
+        foreach ($this->failPointsAndServers as [$failPoint, $server]) {
+            $operation = new DatabaseCommand('admin', ['configureFailPoint' => $failPoint, 'mode' => 'off']);
+            $operation->execute($server);
+        }
+
+        $this->failPointsAndServers = [];
+    }
+}

tests/UnifiedSpecTests/Operation.php

@@ -16,7 +16,6 @@
 use MongoDB\Model\CollectionInfo;
 use MongoDB\Model\DatabaseInfo;
 use MongoDB\Model\IndexInfo;
-use MongoDB\Operation\DatabaseCommand;
 use MongoDB\Operation\FindOneAndReplace;
 use MongoDB\Operation\FindOneAndUpdate;
 use PHPUnit\Framework\Assert;
@@ -909,15 +908,17 @@ private function executeForTestRunner()
                 assertArrayHasKey('failPoint', $args);
                 assertInstanceOf(Client::class, $args['client']);
                 assertInstanceOf(stdClass::class, $args['failPoint']);
-                $this->createFailPoint($args['failPoint'], $args['client']->getManager()->selectServer());
+                // Configure the fail point via the Context so it can later be disabled
+                $this->context->configureFailPoint($args['failPoint'], $args['client']->getManager()->selectServer());
                 break;
             case 'targetedFailPoint':
                 assertArrayHasKey('session', $args);
                 assertArrayHasKey('failPoint', $args);
                 assertInstanceOf(Session::class, $args['session']);
                 assertInstanceOf(stdClass::class, $args['failPoint']);
                 assertNotNull($args['session']->getServer(), 'Session is pinned');
-                $this->createFailPoint($args['failPoint'], $args['session']->getServer());
+                // Configure the fail point via the Context so it can later be disabled
+                $this->context->configureFailPoint($args['failPoint'], $args['session']->getServer());
                 break;
             case 'loop':
                 assertArrayHasKey('operations', $args);
@@ -936,14 +937,6 @@ private function executeForTestRunner()
         }
     }
 
-    private function createFailPoint(stdClass $failPoint, Server $server): void
-    {
-        $operation = new DatabaseCommand('admin', $failPoint);
-        $operation->execute($server);
-
-        $this->context->registerFailPoint($failPoint, $server);
-    }
-
     private function getIndexNames(string $databaseName, string $collectionName): array
     {
         return array_map(

tests/UnifiedSpecTests/UnifiedTestRunner.php

@@ -217,7 +217,7 @@ private function doTestCase(stdClass $test, string $schemaVersion, ?array $runOn
             $context->stopEventObservers();
             $context->stopEventCollectors();
         } finally {
-            $this->disableFailPoints($context->getFailPoints());
+            $context->disableFailPoints();
         }
 
         if (isset($test->expectEvents)) {
@@ -231,15 +231,6 @@ private function doTestCase(stdClass $test, string $schemaVersion, ?array $runOn
         }
     }
 
-    /** @param list<array{failPoint: stdClass, server: Server}> $failPoints */
-    private function disableFailPoints(array $failPoints): void
-    {
-        foreach ($failPoints as ['failPoint' => $failPoint, 'server' => $server]) {
-            $operation = new DatabaseCommand('admin', ['configureFailPoint' => $failPoint, 'mode' => 'off']);
-            $operation->execute($server);
-        }
-    }
-
     /**
      * Checks server version and topology requirements.
      *