diff --git a/core/lib/Drupal/Core/Entity/Sql/TableMappingInterface.php b/core/lib/Drupal/Core/Entity/Sql/TableMappingInterface.php
index caca02d..1107122 100644
--- a/core/lib/Drupal/Core/Entity/Sql/TableMappingInterface.php
+++ b/core/lib/Drupal/Core/Entity/Sql/TableMappingInterface.php
@@ -50,14 +50,14 @@ public function getAllColumns($table_name);
    *   array(
    *     'description' => array(
    *       'value' => 'description__value',
-   *       'format' => 'description_format',
+   *       'format' => 'description__format',
    *     ),
    *   ),
    * @endcode
    *
-   * The values of the inner array ('title', 'description__value', and
-   * 'description_format' in the example above) depend on the table mapping
-   * implementation.
+   * The values of the inner array ('title' for the single-column field and
+   * 'description__value', and 'description__format' for the multi-column field
+   * in the example above) depend on the table mapping implementation.
    *
    * @param string $table_name
    *   The name of the table to return the columns for.
diff --git a/core/tests/Drupal/Tests/Core/Entity/Schema/ContentEntitySchemaHandlerTest.php b/core/tests/Drupal/Tests/Core/Entity/Schema/ContentEntitySchemaHandlerTest.php
index 1b1891d..390fc57 100644
--- a/core/tests/Drupal/Tests/Core/Entity/Schema/ContentEntitySchemaHandlerTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/Schema/ContentEntitySchemaHandlerTest.php
@@ -111,17 +111,26 @@ public function testGetSchemaLayoutBase() {
 
     $this->setupSchemaHandler();
 
+    $table_mapping = $this->getMock('Drupal\Core\Entity\Sql\TableMappingInterface');
+    $table_mapping->expects($this->once())
+      ->method('getTableNames')
+      ->will($this->returnValue(array('entity_test')));
+    $table_mapping->expects($this->once())
+      ->method('getFieldColumns')
+      ->with('entity_test')
+      ->will($this->returnValue(array(
+        'id' => array('value' => 'id'),
+        'name' => array('value' => 'name'),
+        'type' => array('value' => 'type'),
+      )));
+    $table_mapping->expects($this->once())
+      ->method('getExtraColumns')
+      ->with('entity_test')
+      ->will($this->returnValue(array()));
+
     $this->storage->expects($this->once())
       ->method('getTableMapping')
-      ->will($this->returnValue(
-        array(
-          'entity_test' => array(
-            'id' => array('value' => 'id'),
-            'name' => array('value' => 'name'),
-            'type' => array('value' => 'type'),
-          ),
-        )
-      ));
+      ->will($this->returnValue($table_mapping));
 
     $expected = array(
       'entity_test' => array(
@@ -187,18 +196,27 @@ public function testGetSchemaLayoutBaseWithUuid() {
 
     $this->setupSchemaHandler();
 
+    $table_mapping = $this->getMock('Drupal\Core\Entity\Sql\TableMappingInterface');
+    $table_mapping->expects($this->once())
+      ->method('getTableNames')
+      ->will($this->returnValue(array('entity_test')));
+    $table_mapping->expects($this->once())
+      ->method('getFieldColumns')
+      ->with('entity_test')
+      ->will($this->returnValue(array(
+        'id' => array('value' => 'id'),
+        'uuid' => array('value' => 'uuid'),
+        'name' => array('value' => 'name'),
+        'type' => array('value' => 'type'),
+      )));
+    $table_mapping->expects($this->once())
+      ->method('getExtraColumns')
+      ->with('entity_test')
+      ->will($this->returnValue(array()));
+
     $this->storage->expects($this->once())
       ->method('getTableMapping')
-      ->will($this->returnValue(
-        array(
-          'entity_test' => array(
-            'id' => array('value' => 'id'),
-            'uuid' => array('value' => 'uuid'),
-            'name' => array('value' => 'name'),
-            'type' => array('value' => 'type'),
-          ),
-        )
-      ));
+      ->will($this->returnValue($table_mapping));
 
     $expected = array(
       'entity_test' => array(