diff --git a/core/lib/Drupal/Core/Entity/ContentEntityDatabaseStorage.php b/core/lib/Drupal/Core/Entity/ContentEntityDatabaseStorage.php
index c28dbf5..25cf590 100644
--- a/core/lib/Drupal/Core/Entity/ContentEntityDatabaseStorage.php
+++ b/core/lib/Drupal/Core/Entity/ContentEntityDatabaseStorage.php
@@ -165,6 +165,7 @@ public function __construct(EntityTypeInterface $entity_type, Connection $databa
   protected function initTableLayout() {
     // Reset table field values to ensure changes in the entity type deifnition
     // are correctly reflected int the table layout.
+    $this->tableMapping = NULL;
     $this->revisionKey = NULL;
     $this->revisionTable = NULL;
     $this->dataTable = NULL;
@@ -258,7 +259,7 @@ protected function schemaHandler() {
   public function setEntityType(ContentEntityTypeInterface $entity_type) {
     if ($this->entityType->id() == $entity_type->id()) {
       $this->entityType = $entity_type;
-      $this->tableMapping = NULL;
+      $this->initTableLayout();
     }
     else {
       throw new EntityStorageException(String::format('Unsupported entity type @id', array('@id' => $entity_type->id())));
diff --git a/core/lib/Drupal/Core/Entity/Schema/ContentEntitySchemaHandler.php b/core/lib/Drupal/Core/Entity/Schema/ContentEntitySchemaHandler.php
index 78b3187..bbc0356 100644
--- a/core/lib/Drupal/Core/Entity/Schema/ContentEntitySchemaHandler.php
+++ b/core/lib/Drupal/Core/Entity/Schema/ContentEntitySchemaHandler.php
@@ -132,7 +132,6 @@ public function updateEntitySchema(ContentEntityTypeInterface $entity_type, Cont
         $this->storage->setEntityType($entity_type);
         $entity_type_id = $entity_type->id();
         unset($this->schema[$entity_type_id]);
-        $this->fieldStorageDefinitions = $this->entityManager->getFieldStorageDefinitions($entity_type_id);
         $this->createEntitySchema($entity_type);
       }
       catch (\Exception $e) {
@@ -174,15 +173,19 @@ protected function getEntitySchema(ContentEntityTypeInterface $entity_type) {
       // We need to act only on entity schema tables.
       $table_mapping = $this->storage->getTableMapping();
       $table_names = array_diff($table_mapping->getTableNames(), $table_mapping->getDedicatedTableNames());
+      $storage_definitions = $this->entityManager->getFieldStorageDefinitions($entity_type_id);
       foreach ($table_names as $table_name) {
         if (!isset($schema[$table_name])) {
           $schema[$table_name] = array();
         }
         foreach ($table_mapping->getFieldNames($table_name) as $field_name) {
+          if (!isset($storage_definitions[$field_name])) {
+            throw new FieldException(String::format('Fieled storage definition for "@field_name" could not be found.', array('@field_name' => $field_name)));
+          }
           // Add the schema for base field definitions.
-          if (isset($this->fieldStorageDefinitions[$field_name]) && $table_mapping->allowsSharedTableStorage($this->fieldStorageDefinitions[$field_name])) {
+          elseif ($table_mapping->allowsSharedTableStorage($storage_definitions[$field_name])) {
             $column_names = $table_mapping->getColumnNames($field_name);
-            $storage_definition = $this->fieldStorageDefinitions[$field_name];
+            $storage_definition = $storage_definitions[$field_name];
             $schema[$table_name] = array_merge_recursive($schema[$table_name], $this->getSharedTableFieldSchema($storage_definition, $column_names));
           }
         }
diff --git a/core/modules/system/src/Tests/Entity/EntitySchemaTest.php b/core/modules/system/src/Tests/Entity/EntitySchemaTest.php
index 22f1a72..61e1b33 100644
--- a/core/modules/system/src/Tests/Entity/EntitySchemaTest.php
+++ b/core/modules/system/src/Tests/Entity/EntitySchemaTest.php
@@ -7,6 +7,8 @@
 
 namespace Drupal\system\Tests\Entity;
 
+use Drupal\Component\Utility\String;
+
 /**
  * Tests adding a custom bundle field.
  */
@@ -60,7 +62,7 @@ public function setUp() {
    */
   public function testCustomFieldCreateDelete() {
     // Install the module which adds the field.
-    $this->installModuleAndRefreshServices('entity_schema_test');
+    $this->installModule('entity_schema_test');
     $this->entityManager->clearCachedDefinitions();
     $definition = $this->entityManager->getBaseFieldDefinitions('entity_test')['custom_base_field'];
     $this->assertNotNull($definition, 'Base field definition found.');
@@ -76,7 +78,7 @@ public function testCustomFieldCreateDelete() {
 
     $table = $table_mapping->getDedicatedDataTableName($definition->getFieldStorageDefinition());
     $this->assertTrue($this->database->schema()->tableExists($table), 'Table created');
-    $this->uninstallModuleAndRefreshServices('entity_schema_test');
+    $this->uninstallModule('entity_schema_test');
     $this->assertFalse($this->database->schema()->fieldExists($base_table, $base_column), 'Table column dropped');
     $this->assertFalse($this->database->schema()->tableExists($table), 'Table dropped');
   }
@@ -85,37 +87,37 @@ public function testCustomFieldCreateDelete() {
    * Tests that entity schema responds to changes in the entity type definition.
    */
   public function testEntitySchemaUpdate() {
-    $this->installModuleAndRefreshServices('entity_schema_test');
+    $this->installModule('entity_schema_test');
     $schema_handler = $this->database->schema();
     $tables = array('entity_test', 'entity_test_revision', 'entity_test_field_data', 'entity_test_field_revision');
     $dedicated_tables = array('entity_test__custom_bundle_field', 'entity_test_revision__custom_bundle_field');
 
     // Initially only the base table should exist.
     foreach ($tables as $index => $table) {
-      $this->assertEqual($schema_handler->tableExists($table), !$index, format_string('Entity schema correct for the @table table.', array('@table' => $table)));
+      $this->assertEqual($schema_handler->tableExists($table), !$index, String::format('Entity schema correct for the @table table.', array('@table' => $table)));
     }
     foreach ($dedicated_tables as $table) {
-      $this->assertTrue($schema_handler->tableExists($table), format_string('Field schema correct for the @table table.', array('@table' => $table)));
+      $this->assertTrue($schema_handler->tableExists($table), String::format('Field schema correct for the @table table.', array('@table' => $table)));
     }
 
     // Update the entity type definition and check that the entity schema now
     // supports translations and revisions.
     $this->updateEntityType(TRUE);
     foreach ($tables as $table) {
-      $this->assertTrue($schema_handler->tableExists($table), format_string('Entity schema correct for the @table table.', array('@table' => $table)));
+      $this->assertTrue($schema_handler->tableExists($table), String::format('Entity schema correct for the @table table.', array('@table' => $table)));
     }
     foreach ($dedicated_tables as $table) {
-      $this->assertTrue($schema_handler->tableExists($table), format_string('Field schema correct for the @table table.', array('@table' => $table)));
+      $this->assertTrue($schema_handler->tableExists($table), String::format('Field schema correct for the @table table.', array('@table' => $table)));
     }
 
     // Revert changes and check that the entity schema now does not support
     // neither translations nor revisions.
     $this->updateEntityType(FALSE);
     foreach ($tables as $index => $table) {
-      $this->assertEqual($schema_handler->tableExists($table), !$index, format_string('Entity schema correct for the @table table.', array('@table' => $table)));
+      $this->assertEqual($schema_handler->tableExists($table), !$index, String::format('Entity schema correct for the @table table.', array('@table' => $table)));
     }
     foreach ($dedicated_tables as $table) {
-      $this->assertTrue($schema_handler->tableExists($table), format_string('Field schema correct for the @table table.', array('@table' => $table)));
+      $this->assertTrue($schema_handler->tableExists($table), String::format('Field schema correct for the @table table.', array('@table' => $table)));
     }
   }
 
@@ -139,13 +141,9 @@ protected function updateEntityType($alter) {
    * @param string $module
    *   The module to install.
    */
-  protected function installModuleAndRefreshServices($module) {
+  protected function installModule($module) {
     $this->moduleHandler->install(array($module), FALSE);
-    $this->container = \Drupal::getContainer();
-    $this->moduleHandler = $this->container->get('module_handler');
-    $this->database = $this->container->get('database');
-    $this->entityManager = $this->container->get('entity.manager');
-    $this->state = $this->container->get('state');
+    $this->refreshServices();
   }
 
   /**
@@ -154,8 +152,15 @@ protected function installModuleAndRefreshServices($module) {
    * @param string $module
    *   The module to uninstall.
    */
-  protected function uninstallModuleAndRefreshServices($module) {
+  protected function uninstallModule($module) {
     $this->moduleHandler->uninstall(array($module), FALSE);
+    $this->refreshServices();
+  }
+
+  /**
+   * Refresh services.
+   */
+  protected function refreshServices() {
     $this->container = \Drupal::getContainer();
     $this->moduleHandler = $this->container->get('module_handler');
     $this->database = $this->container->get('database');