diff --git a/core/modules/entity/tests/modules/entity_test/lib/Drupal/entity_test/EntityTestStorageController.php b/core/modules/entity/tests/modules/entity_test/lib/Drupal/entity_test/EntityTestStorageController.php
index 00e9cf5..ad53689 100644
--- a/core/modules/entity/tests/modules/entity_test/lib/Drupal/entity_test/EntityTestStorageController.php
+++ b/core/modules/entity/tests/modules/entity_test/lib/Drupal/entity_test/EntityTestStorageController.php
@@ -45,11 +45,11 @@ class EntityTestStorageController extends DatabaseStorageController {
         unset($values['default_langcode']);
       }
 
-      $entity_fields = $this->entityInfo['schema_fields_sql']['base table'];
+      $data_schema = drupal_get_schema('entity_test_property_data');
       $query->addField('data', $this->idKey);
       foreach ($values as $field => $value) {
         // Check on which table the condition needs to be added.
-        $table = in_array($field, $entity_fields) ? 'base' : 'data';
+        $table = isset($data_schema['fields'][$field]) ? 'data' : 'base';
         $query->condition($table . '.' . $field, $value);
       }
     }