Make NameTableStore use LoadBalancer::getConnectionRef()

[lhc/web/wiklou.git] / tests / phpunit / includes / Storage / NameTableStoreTest.php

diff --git a/tests/phpunit/includes/Storage/NameTableStoreTest.php b/tests/phpunit/includes/Storage/NameTableStoreTest.php
index ca87b49..7a4ea2d 100644 (file)
--- a/tests/phpunit/includes/Storage/NameTableStoreTest.php
+++ b/tests/phpunit/includes/Storage/NameTableStoreTest.php
@@ -10,8 +10,9 @@ use MediaWiki\Storage\NameTableStore;
 use MediaWikiTestCase;
 use Psr\Log\NullLogger;
 use WANObjectCache;
-use Wikimedia\Rdbms\Database;
+use Wikimedia\Rdbms\IDatabase;
 use Wikimedia\Rdbms\LoadBalancer;
+use Wikimedia\Rdbms\MaintainableDBConnRef;
 use Wikimedia\TestingAccessWrapper;
 
 /**
@@ -51,43 +52,33 @@ class NameTableStoreTest extends MediaWikiTestCase {
                        ->disableOriginalConstructor()
                        ->getMock();
                $mock->expects( $this->any() )
-                       ->method( 'getConnection' )
-                       ->willReturn( $db );
+                       ->method( 'getConnectionRef' )
+                       ->willReturnCallback( function ( $i ) use ( $mock, $db ) {
+                               return new MaintainableDBConnRef( $mock, $db, $i );
+                       } );
                return $mock;
        }
 
        private function getCallCheckingDb( $insertCalls, $selectCalls ) {
-               $mock = $this->getMockBuilder( Database::class )
+               $proxiedMethods = [
+                       'select' => $selectCalls,
+                       'insert' => $insertCalls,
+                       'affectedRows' => null,
+                       'insertId' => null,
+                       'getSessionLagStatus' => null,
+                       'writesPending' => null,
+                       'onTransactionPreCommitOrIdle' => null
+               ];
+               $mock = $this->getMockBuilder( IDatabase::class )
                        ->disableOriginalConstructor()
                        ->getMock();
-               $mock->expects( $this->exactly( $insertCalls ) )
-                       ->method( 'insert' )
-                       ->willReturnCallback( function ( ...$args ) {
-                               return call_user_func_array( [ $this->db, 'insert' ], $args );
-                       } );
-               $mock->expects( $this->exactly( $selectCalls ) )
-                       ->method( 'select' )
-                       ->willReturnCallback( function ( ...$args ) {
-                               return call_user_func_array( [ $this->db, 'select' ], $args );
-                       } );
-               $mock->expects( $this->exactly( $insertCalls ) )
-                       ->method( 'affectedRows' )
-                       ->willReturnCallback( function ( ...$args ) {
-                               return call_user_func_array( [ $this->db, 'affectedRows' ], $args );
-                       } );
-               $mock->expects( $this->any() )
-                       ->method( 'insertId' )
-                       ->willReturnCallback( function ( ...$args ) {
-                               return call_user_func_array( [ $this->db, 'insertId' ], $args );
-                       } );
-               $mock->expects( $this->any() )
-                       ->method( 'query' )
-                       ->willReturn( [] );
-               $mock->expects( $this->any() )
-                       ->method( 'isOpen' )
-                       ->willReturn( true );
-               $wrapper = TestingAccessWrapper::newFromObject( $mock );
-               $wrapper->queryLogger = new NullLogger();
+               foreach ( $proxiedMethods as $method => $count ) {
+                       $mock->expects( is_int( $count ) ? $this->exactly( $count ) : $this->any() )
+                               ->method( $method )
+                               ->willReturnCallback( function ( ...$args ) use ( $method ) {
+                                       return call_user_func_array( [ $this->db, $method ], $args );
+                               } );
+               }
                return $mock;
        }