diff --git a/core/tests/Drupal/Tests/Core/Entity/TypedData/EntityAdapterUnitTest.php b/core/tests/Drupal/Tests/Core/Entity/TypedData/EntityAdapterUnitTest.php
index bb97044..202aef4 100644
--- a/core/tests/Drupal/Tests/Core/Entity/TypedData/EntityAdapterUnitTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/TypedData/EntityAdapterUnitTest.php
@@ -148,8 +148,15 @@ protected function setUp() {
       ->getMock();
     $this->typedDataManager->expects($this->any())
       ->method('getDefinition')
-      ->with('entity')
-      ->will($this->returnValue(['class' => '\Drupal\Core\Entity\Plugin\DataType\EntityAdapter']));
+      ->will($this->returnCallback(function($plugin_id){
+        switch ($plugin_id) {
+          case 'entity':
+            return ['class' => '\Drupal\Core\Entity\Plugin\DataType\EntityAdapter'];
+          default:
+            // @todo: maybe use a reasonable argument value
+            return ['list_class' => '\Drupal\Core\Field\FieldItemList'];
+        }
+      }));
     $this->typedDataManager->expects($this->any())
       ->method('getDefaultConstraints')
       ->willReturn([]);
@@ -190,20 +197,12 @@ protected function setUp() {
       ->method('createFieldItemList')
       ->willReturn($this->fieldItemList);
 
-    $typed_data_manager = $this->getMock('Drupal\Component\Plugin\PluginManagerInterface');
-    // @todo: maybe use a reasonable argument
-    $typed_data_manager->expects($this->any())
-      ->method('getDefinition')
-      ->with($this->anything())
-      ->will($this->returnValue(['list_class' => '\Drupal\Core\Field\FieldItemList']));
-
     $container = new ContainerBuilder();
     $container->set('entity.manager', $this->entityManager);
     $container->set('uuid', $this->uuid);
     $container->set('typed_data_manager', $this->typedDataManager);
     $container->set('language_manager', $this->languageManager);
     $container->set('plugin.manager.field.field_type', $this->fieldTypePluginManager);
-    $container->set('typed_data_manager', $typed_data_manager);
     \Drupal::setContainer($container);
 
     $this->fieldDefinitions = array(