Allow schema overrides to drop tables.

This allows unit tests to use schema overrides that drop tables,
in addition to overrides that create or modify tables.

Change-Id: I59761c7db7f83698749324ca6b9ffced86ab1249

4 files changed, 148 insertions, 61 deletions

diff --git a/tests/phpunit/MediaWikiTestCase.php b/tests/phpunit/MediaWikiTestCase.php
index 92c0714eef4d..0d2b788b221e 100644
--- a/tests/phpunit/MediaWikiTestCase.php
+++ b/tests/phpunit/MediaWikiTestCase.php
@@ -1313,57 +1313,113 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 		}
 	}
 
+	private static $schemaOverrideDefaults = [
+		'scripts' => [],
+		'create' => [],
+		'drop' => [],
+		'alter' => [],
+	];
+
 	/**
 	 * Stub. If a test suite needs to test against a specific database schema, it should
 	 * override this method and return the appropriate information from it.
 	 *
-	 * @return [ $tables, $scripts ] A tuple of two lists, with $tables being a list of tables
-	 *         that will be re-created by the scripts, and $scripts being a list of SQL script
-	 *         files for creating the tables listed.
+	 * @param IMaintainableDatabase $db The DB connection to use for the mock schema.
+	 *        May be used to check the current state of the schema, to determine what
+	 *        overrides are needed.
+	 *
+	 * @return array An associative array with the following fields:
+	 *  - 'scripts': any SQL scripts to run. If empty or not present, schema overrides are skipped.
+	 * - 'create': A list of tables created (may or may not exist in the original schema).
+	 * - 'drop': A list of tables dropped (expected to be present in the original schema).
+	 * - 'alter': A list of tables altered (expected to be present in the original schema).
 	 */
-	protected function getSchemaOverrides() {
-		return [ [], [] ];
+	protected function getSchemaOverrides( IMaintainableDatabase $db ) {
+		return [];
+	}
+
+	/**
+	 * Undoes the dpecified schema overrides..
+	 * Called once per test class, just before addDataOnce().
+	 *
+	 * @param IMaintainableDatabase $db
+	 * @param array $oldOverrides
+	 */
+	private function undoSchemaOverrides( IMaintainableDatabase $db, $oldOverrides ) {
+		$this->ensureMockDatabaseConnection( $db );
+
+		$oldOverrides = $oldOverrides + self::$schemaOverrideDefaults;
+		$originalTables = $this->listOriginalTables( $db );
+
+		// Drop tables that need to be restored or removed.
+		$tablesToDrop = array_merge( $oldOverrides['create'], $oldOverrides['alter'] );
+
+		// Restore tables that have been dropped or created or altered,
+		// if they exist in the original schema.
+		$tablesToRestore = array_merge( $tablesToDrop, $oldOverrides['drop'] );
+		$tablesToRestore = array_intersect( $originalTables, $tablesToRestore );
+
+		if ( $tablesToDrop ) {
+			$this->dropMockTables( $db, $tablesToDrop );
+		}
+
+		if ( $tablesToRestore ) {
+			$this->recloneMockTables( $db, $tablesToRestore );
+		}
 	}
 
 	/**
-	 * Applies any schema changes requested by calling setDbSchema().
+	 * Applies the schema overrides returned by getSchemaOverrides(),
+	 * after undoing any previously applied schema overrides.
 	 * Called once per test class, just before addDataOnce().
 	 */
 	private function setUpSchema( IMaintainableDatabase $db ) {
-		list( $tablesToAlter, $scriptsToRun ) = $this->getSchemaOverrides();
+		// Undo any active overrides.
+		$oldOverrides = isset( $db->_schemaOverrides ) ? $db->_schemaOverrides
+			: self::$schemaOverrideDefaults;
+
+		if ( $oldOverrides['alter'] || $oldOverrides['create'] || $oldOverrides['drop'] ) {
+			$this->undoSchemaOverrides( $db, $oldOverrides );
+		}
+
+		// Determine new overrides.
+		$overrides = $this->getSchemaOverrides( $db ) + self::$schemaOverrideDefaults;
+
+		$extraKeys = array_diff(
+			array_keys( $overrides ),
+			array_keys( self::$schemaOverrideDefaults )
+		);
 
-		if ( $tablesToAlter && !$scriptsToRun ) {
+		if ( $extraKeys ) {
 			throw new InvalidArgumentException(
-				'No scripts supplied for applying the database schema.'
+				'Schema override contains extra keys: ' . var_export( $extraKeys, true )
 			);
 		}
 
-		if ( !$tablesToAlter && $scriptsToRun ) {
+		if ( !$overrides['scripts'] ) {
+			// no scripts to run
+			return;
+		}
+
+		if ( !$overrides['create'] && !$overrides['drop'] && !$overrides['alter'] ) {
 			throw new InvalidArgumentException(
-				'No tables declared to be altered by schema scripts.'
+				'Schema override scripts given, but no tables are declared to be '
+				. 'created, dropped or altered.'
 			);
 		}
 
 		$this->ensureMockDatabaseConnection( $db );
 
-		$previouslyAlteredTables = isset( $db->_alteredMockTables ) ? $db->_alteredMockTables : [];
-
-		if ( !$tablesToAlter && !$previouslyAlteredTables ) {
-			return; // nothing to do
-		}
-
-		$tablesToDrop = array_merge( $previouslyAlteredTables, $tablesToAlter );
-		$tablesToRestore = array_diff( $previouslyAlteredTables, $tablesToAlter );
+		// Drop the tables that will be created by the schema scripts.
+		$originalTables = $this->listOriginalTables( $db );
+		$tablesToDrop = array_intersect( $originalTables, $overrides['create'] );
 
 		if ( $tablesToDrop ) {
 			$this->dropMockTables( $db, $tablesToDrop );
 		}
 
-		if ( $tablesToRestore ) {
-			$this->recloneMockTables( $db, $tablesToRestore );
-		}
-
-		foreach ( $scriptsToRun as $script ) {
+		// Run schema override scripts.
+		foreach ( $overrides['scripts'] as $script ) {
 			$db->sourceFile(
 				$script,
 				null,
@@ -1375,7 +1431,7 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 			);
 		}
 
-		$db->_alteredMockTables = $tablesToAlter;
+		$db->_schemaOverrides = $overrides;
 	}
 
 	private function mungeSchemaUpdateQuery( $cmd ) {
@@ -1406,7 +1462,24 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 	}
 
 	/**
+	 * Lists all tables in the live database schema.
+	 *
+	 * @param IMaintainableDatabase $db
+	 * @return array
+	 */
+	private function listOriginalTables( IMaintainableDatabase $db ) {
+		if ( !isset( $db->_originalTablePrefix ) ) {
+			throw new LogicException( 'No original table prefix know, cannot list tables!' );
+		}
+
+		$originalTables = $db->listTables( $db->_originalTablePrefix, __METHOD__ );
+		return $originalTables;
+	}
+
+	/**
 	 * Re-clones the given mock tables to restore them based on the live database schema.
+	 * The tables listed in $tables are expected to currently not exist, so dropMockTables()
+	 * should be called first.
 	 *
 	 * @param IMaintainableDatabase $db
 	 * @param array $tables
@@ -1418,7 +1491,7 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 			throw new LogicException( 'No original table prefix know, cannot restore tables!' );
 		}
 
-		$originalTables = $db->listTables( $db->_originalTablePrefix, __METHOD__ );
+		$originalTables = $this->listOriginalTables( $db );
 		$tables = array_intersect( $tables, $originalTables );
 
 		$dbClone = new CloneDatabase( $db, $tables, $db->tablePrefix(), $db->_originalTablePrefix );
@@ -1456,6 +1529,10 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 					continue;
 				}
 
+				if ( !$db->tableExists( $tbl ) ) {
+					continue;
+				}
+
 				if ( $truncate ) {
 					$db->query( 'TRUNCATE TABLE ' . $db->tableName( $tbl ), __METHOD__ );
 				} else {
diff --git a/tests/phpunit/tests/MediaWikiTestCaseSchema1Test.php b/tests/phpunit/tests/MediaWikiTestCaseSchema1Test.php
index 6f944946faf0..d794d1318af8 100644
--- a/tests/phpunit/tests/MediaWikiTestCaseSchema1Test.php
+++ b/tests/phpunit/tests/MediaWikiTestCaseSchema1Test.php
@@ -1,4 +1,5 @@
 <?php
+use Wikimedia\Rdbms\IMaintainableDatabase;
 
 /**
  * @covers MediaWikiTestCase
@@ -10,10 +11,12 @@ class MediaWikiTestCaseSchema1Test extends MediaWikiTestCase {
 
 	public static $hasRun = false;
 
-	public function getSchemaOverrides() {
+	public function getSchemaOverrides( IMaintainableDatabase $db ) {
 		return [
-			[ 'imagelinks', 'MediaWikiTestCaseTestTable' ],
-			[ __DIR__ . '/MediaWikiTestCaseSchemaTest.sql' ]
+			'create' => [ 'MediaWikiTestCaseTestTable', 'imagelinks' ],
+			'drop' => [ 'oldimage' ],
+			'alter' => [ 'pagelinks' ],
+			'scripts' => [ __DIR__ . '/MediaWikiTestCaseSchemaTest.sql' ]
 		];
 	}
 
@@ -23,37 +26,26 @@ class MediaWikiTestCaseSchema1Test extends MediaWikiTestCase {
 		$this->assertTrue( self::$hasRun );
 	}
 
-	public function testSchemaExtension() {
-		// make sure we can use the MediaWikiTestCaseTestTable table
-
-		$input = [ 'id' => '5', 'name' => 'Test' ];
-
-		$this->db->insert(
-			'MediaWikiTestCaseTestTable',
-			$input
-		);
-
-		$output = $this->db->selectRow( 'MediaWikiTestCaseTestTable', array_keys( $input ), [] );
-		$this->assertEquals( (object)$input, $output );
+	public function testTableWasCreated() {
+		// Make sure MediaWikiTestCaseTestTable was created.
+		$this->assertTrue( $this->db->tableExists( 'MediaWikiTestCaseTestTable' ) );
 	}
 
-	public function testSchemaOverride() {
-		// make sure we can use the il_frobniz field
-
-		$input = [
-			'il_from' => '7',
-			'il_from_namespace' => '0',
-			'il_to' => 'Foo.jpg',
-			'il_frobniz' => 'Xyzzy',
-		];
+	public function testTableWasDropped() {
+		// Make sure oldimage was dropped
+		$this->assertFalse( $this->db->tableExists( 'oldimage' ) );
+	}
 
-		$this->db->insert(
-			'imagelinks',
-			$input
-		);
+	public function testTableWasOverriden() {
+		// Make sure imagelinks was overwritten
+		$this->assertTrue( $this->db->tableExists( 'imagelinks' ) );
+		$this->assertTrue( $this->db->fieldExists( 'imagelinks', 'il_frobnitz' ) );
+	}
 
-		$output = $this->db->selectRow( 'imagelinks', array_keys( $input ), [] );
-		$this->assertEquals( (object)$input, $output );
+	public function testTableWasAltered() {
+		// Make sure pagelinks was altered
+		$this->assertTrue( $this->db->tableExists( 'pagelinks' ) );
+		$this->assertTrue( $this->db->fieldExists( 'pagelinks', 'pl_frobnitz' ) );
 	}
 
 }
diff --git a/tests/phpunit/tests/MediaWikiTestCaseSchema2Test.php b/tests/phpunit/tests/MediaWikiTestCaseSchema2Test.php
index 74f053eb3ed1..5464dc43f814 100644
--- a/tests/phpunit/tests/MediaWikiTestCaseSchema2Test.php
+++ b/tests/phpunit/tests/MediaWikiTestCaseSchema2Test.php
@@ -19,17 +19,30 @@ class MediaWikiTestCaseSchema2Test extends MediaWikiTestCase {
 		$this->assertTrue( MediaWikiTestCaseSchema1Test::$hasRun );
 	}
 
-	public function testSchemaExtension() {
+	public function testCreatedTableWasRemoved() {
 		// Make sure MediaWikiTestCaseTestTable created by MediaWikiTestCaseSchema1Test
 		// was dropped before executing MediaWikiTestCaseSchema2Test.
 		$this->assertFalse( $this->db->tableExists( 'MediaWikiTestCaseTestTable' ) );
 	}
 
-	public function testSchemaOverride() {
-		// Make sure imagelinks modified by MediaWikiTestCaseSchema1Test
+	public function testDroppedTableWasRestored() {
+		// Make sure oldimage that was dropped by MediaWikiTestCaseSchema1Test
+		// was restored before executing MediaWikiTestCaseSchema2Test.
+		$this->assertTrue( $this->db->tableExists( 'oldimage' ) );
+	}
+
+	public function testOverridenTableWasRestored() {
+		// Make sure imagelinks overwritten by MediaWikiTestCaseSchema1Test
 		// was restored to the original schema before executing MediaWikiTestCaseSchema2Test.
 		$this->assertTrue( $this->db->tableExists( 'imagelinks' ) );
-		$this->assertFalse( $this->db->fieldExists( 'imagelinks', 'il_frobniz' ) );
+		$this->assertFalse( $this->db->fieldExists( 'imagelinks', 'il_frobnitz' ) );
+	}
+
+	public function testAlteredTableWasRestored() {
+		// Make sure pagelinks altered by MediaWikiTestCaseSchema1Test
+		// was restored to the original schema before executing MediaWikiTestCaseSchema2Test.
+		$this->assertTrue( $this->db->tableExists( 'pagelinks' ) );
+		$this->assertFalse( $this->db->fieldExists( 'pagelinks', 'pl_frobnitz' ) );
 	}
 
 }
diff --git a/tests/phpunit/tests/MediaWikiTestCaseSchemaTest.sql b/tests/phpunit/tests/MediaWikiTestCaseSchemaTest.sql
index 58460e2e9fea..e2818b552aee 100644
--- a/tests/phpunit/tests/MediaWikiTestCaseSchemaTest.sql
+++ b/tests/phpunit/tests/MediaWikiTestCaseSchemaTest.sql
@@ -8,6 +8,11 @@ CREATE TABLE /*_*/imagelinks (
   il_from int NOT NULL DEFAULT 0,
   il_from_namespace int NOT NULL DEFAULT 0,
   il_to varchar(127) NOT NULL DEFAULT '',
-  il_frobniz varchar(127) NOT NULL DEFAULT 'FROB',
+  il_frobnitz varchar(127) NOT NULL DEFAULT 'FROB',
   PRIMARY KEY (il_from,il_to)
 ) /*$wgDBTableOptions*/;
+
+ALTER TABLE /*_*/pagelinks
+ADD pl_frobnitz varchar(127) NOT NULL DEFAULT 'FROB';
+
+DROP TABLE /*_*/oldimage;