Return promise from StatementPool::prepare(); update tests

src/StatementPool.php

@@ -27,14 +27,14 @@ abstract class StatementPool implements Statement
     private $prepare;
 
     /**
-     * Performs any necessary actions to the statement to prepare it for execution, returning the same or a new
+     * Performs any necessary actions to the statement to prepare it for execution, returning a promise for the same or
-     * Statement object if necessary.
+     * a new Statement object if necessary.
      *
      * @param Statement $statement
      *
-     * @return Statement
+     * @return Promise<Statement>
      */
-    abstract protected function prepare(Statement $statement): Statement;
+    abstract protected function prepare(Statement $statement): Promise;
 
     /**
      * @param ResultSet $resultSet
@@ -97,7 +97,8 @@ abstract class StatementPool implements Statement
             \assert($statement instanceof Statement);
 
             try {
-                $statement = $this->prepare($statement);
+                $statement = yield $this->prepare($statement);
+                \assert($statement instanceof Statement);
                 $result = yield $statement->execute($params);
             } catch (\Throwable $exception) {
                 $this->push($statement);

test/StatementPoolTest.php

@@ -8,6 +8,7 @@ use Amp\PHPUnit\TestCase;
 use Amp\Sql\Pool;
 use Amp\Sql\Statement;
 use Amp\Sql\StatementPool;
+use Amp\Success;
 
 class StatementPoolTest extends TestCase
 {
@@ -35,6 +36,11 @@ class StatementPoolTest extends TestCase
                 ->setConstructorArgs([$pool, $statement, $this->createCallback(0)])
                 ->getMockForAbstractClass();
 
+            $statementPool->method('prepare')
+                ->willReturnCallback(function (Statement $statement) {
+                    return new Success($statement);
+                });
+
             $this->assertTrue($statementPool->isAlive());
             $this->assertSame(\time(), $statementPool->getLastUsedAt());
 
@@ -71,6 +77,11 @@ class StatementPoolTest extends TestCase
                 ->setConstructorArgs([$pool, $statement, $this->createCallback(1)])
                 ->getMockForAbstractClass();
 
+            $statementPool->method('prepare')
+                ->willReturnCallback(function (Statement $statement) {
+                    return new Success($statement);
+                });
+
             $this->assertTrue($statementPool->isAlive());
             $this->assertSame(\time(), $statementPool->getLastUsedAt());