diff --git a/core/modules/simpletest/lib/Drupal/simpletest/DrupalUnitTestBase.php b/core/modules/simpletest/lib/Drupal/simpletest/DrupalUnitTestBase.php
index 64386e9..4af880d 100644
--- a/core/modules/simpletest/lib/Drupal/simpletest/DrupalUnitTestBase.php
+++ b/core/modules/simpletest/lib/Drupal/simpletest/DrupalUnitTestBase.php
@@ -247,9 +247,8 @@ public function containerBuild(ContainerBuilder $container) {
 
       $container
         ->register('keyvalue', 'Drupal\Core\KeyValueStore\KeyValueFactory')
+        ->addMethodCall('setContainer', new Reference('service_container'))
         ->addArgument(new Reference('settings'));
-      $container->set('service_container', newReference('service_container'));
-
       $container->register('state', 'Drupal\Core\KeyValueStore\State')
         ->addArgument(new Reference('keyvalue'));
     }
diff --git a/core/modules/system/lib/Drupal/system/Tests/KeyValueStore/StorageTestBase.php b/core/modules/system/lib/Drupal/system/Tests/KeyValueStore/StorageTestBase.php
index 4b48cfb..ef13704 100644
--- a/core/modules/system/lib/Drupal/system/Tests/KeyValueStore/StorageTestBase.php
+++ b/core/modules/system/lib/Drupal/system/Tests/KeyValueStore/StorageTestBase.php
@@ -57,11 +57,11 @@ protected function setUp() {
       ->setFactoryMethod('getInstance');
     $this->container
       ->register('keyvalue', 'Drupal\Core\KeyValueStore\KeyValueFactory')
-      ->addArgument(new Reference('service_container'))
+      ->addMethodCall('setContainer', new Reference('service_container'))
       ->addArgument(new Reference('settings'));
     $this->container
       ->register('keyvalue.expirable', 'Drupal\Core\KeyValueStore\KeyValueExpirableFactory')
-      ->addArgument(new Reference('service_container'))
+      ->addMethodCall('setContainer', new Reference('service_container'))
       ->addArgument(new Reference('settings'));
     // Define two data collections,
     $this->collections = array(0 => 'zero', 1 => 'one');