diff --git a/page_manager.services.yml b/page_manager.services.yml
index a7d3557..2cdfa8b 100644
--- a/page_manager.services.yml
+++ b/page_manager.services.yml
@@ -9,6 +9,12 @@ services:
     arguments: ['@router.route_provider', '@request_stack']
     tags:
       - { name: 'page_manager_context_provider' }
+  page_manager.static_context_provider:
+    class: Drupal\page_manager\ContextProvider\StaticContextProvider
+    arguments: ['@page_manager.context_mapper']
+    tags:
+      - { name: 'page_manager_context_provider' }
+
   page_manager.context_mapper:
     class: Drupal\page_manager\ContextMapper
     arguments: ['@entity.repository']
diff --git a/page_manager_ui/src/Form/PageVariantEditForm.php b/page_manager_ui/src/Form/PageVariantEditForm.php
index ac25e10..edb5576 100644
--- a/page_manager_ui/src/Form/PageVariantEditForm.php
+++ b/page_manager_ui/src/Form/PageVariantEditForm.php
@@ -351,13 +351,13 @@ protected function buildContextForm() {
 
       // Add operation links if the context is a static context.
       $operations = [];
-      if ($page_variant->getStaticContext($name)) {
+      if ($static_context_name = $this->deriveStaticContextName($name)) {
         $operations['edit'] = [
           'title' => $this->t('Edit'),
           'url' => Url::fromRoute('page_manager.static_context_edit', [
             'page' => $page_variant->get('page'),
             'page_variant' => $page_variant->id(),
-            'name' => $name,
+            'name' => $static_context_name,
           ]),
           'attributes' => $attributes,
         ];
@@ -366,7 +366,7 @@ protected function buildContextForm() {
           'url' => Url::fromRoute('page_manager.static_context_delete', [
             'page' => $page_variant->get('page'),
             'page_variant' => $page_variant->id(),
-            'name' => $name,
+            'name' => $static_context_name,
           ]),
           'attributes' => $attributes,
         ];
@@ -379,7 +379,16 @@ protected function buildContextForm() {
       $form['available_context'][$name] = $row;
     }
     return $form;
+  }
 
+  /**
+   * @todo.
+   */
+  protected function deriveStaticContextName($name) {
+    $static_context_prefix = '@page_manager.static_context_provider:';
+    if (strpos($name, $static_context_prefix) === 0) {
+      return substr($name, strlen($static_context_prefix));
+    }
   }
 
   /**
diff --git a/src/Context/PageManagerContextRepository.php b/src/Context/PageManagerContextRepository.php
index c1004c5..d5708cd 100644
--- a/src/Context/PageManagerContextRepository.php
+++ b/src/Context/PageManagerContextRepository.php
@@ -8,8 +8,8 @@
 namespace Drupal\page_manager\Context;
 
 use Drupal\Core\Plugin\Context\LazyContextRepository;
-use Drupal\page_manager\ContextProvider\PageContextProviderInterface;
-use Drupal\page_manager\PageInterface;
+use Drupal\page_manager\ContextProvider\PageVariantContextProviderInterface;
+use Drupal\page_manager\PageVariantInterface;
 
 /**
  * Provides a context repository which takes a page into account.
@@ -22,7 +22,7 @@ class PageManagerContextRepository extends LazyContextRepository {
   /**
    * {@inheritdoc}
    */
-  public function getRuntimeContexts(array $context_ids, PageInterface $page = NULL) {
+  public function getRuntimeContexts(array $context_ids, PageVariantInterface $page_variant = NULL) {
     $contexts = [];
 
     // Create a map of context providers (service IDs) to unqualified context
@@ -52,8 +52,8 @@ public function getRuntimeContexts(array $context_ids, PageInterface $page = NUL
     // runtime contexts and assign them as requested.
     foreach ($context_ids_by_service as $service_id => $unqualified_context_ids) {
       $service = $this->container->get($service_id);
-      if ($service instanceof PageContextProviderInterface) {
-        $contexts_by_service = $service->getRuntimeContexts($unqualified_context_ids, $page);
+      if ($service instanceof PageVariantContextProviderInterface) {
+        $contexts_by_service = $service->getRuntimeContexts($unqualified_context_ids, $page_variant);
       }
       else {
         $contexts_by_service = $service->getRuntimeContexts($unqualified_context_ids);
@@ -71,12 +71,12 @@ public function getRuntimeContexts(array $context_ids, PageInterface $page = NUL
   /**
    * {@inheritdoc}
    */
-  public function getAvailableContexts(PageInterface $page = NULL) {
+  public function getAvailableContexts(PageVariantInterface $page_variant = NULL) {
     $contexts = [];
     foreach ($this->contextProviderServiceIDs as $service_id) {
       $service = $this->container->get($service_id);
-      if ($service instanceof PageContextProviderInterface) {
-        $contexts_by_service = $service->getAvailableContexts($page);
+      if ($service instanceof PageVariantContextProviderInterface) {
+        $contexts_by_service = $service->getAvailableContexts($page_variant);
       }
       else {
         $contexts_by_service = $service->getAvailableContexts();
diff --git a/src/ContextProvider/PageContextProviderInterface.php b/src/ContextProvider/PageContextProviderInterface.php
deleted file mode 100644
index dde0f68..0000000
--- a/src/ContextProvider/PageContextProviderInterface.php
+++ /dev/null
@@ -1,34 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\page_manager\ContextProvider\PageContextProviderInterface.
- */
-
-namespace Drupal\page_manager\ContextProvider;
-
-use Drupal\Core\Plugin\Context\ContextProviderInterface;
-use Drupal\page_manager\PageInterface;
-
-/**
- * Defines an interface for providing plugin contexts that require a page.
- */
-interface PageContextProviderInterface extends ContextProviderInterface {
-
-  /**
-   * {@inheritdoc}
-   *
-   * @param \Drupal\page_manager\PageInterface $page
-   *   The page to get contexts for.
-   */
-  public function getRuntimeContexts(array $unqualified_context_ids, PageInterface $page = NULL);
-
-  /**
-   * {@inheritdoc}
-   *
-   * @param \Drupal\page_manager\PageInterface $page
-   *   The page to get contexts for.
-   */
-  public function getAvailableContexts(PageInterface $page = NULL);
-
-}
diff --git a/src/ContextProvider/PageVariantContextProviderInterface.php b/src/ContextProvider/PageVariantContextProviderInterface.php
new file mode 100644
index 0000000..3bd8f50
--- /dev/null
+++ b/src/ContextProvider/PageVariantContextProviderInterface.php
@@ -0,0 +1,34 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\page_manager\ContextProvider\PageVariantContextProviderInterface.
+ */
+
+namespace Drupal\page_manager\ContextProvider;
+
+use Drupal\Core\Plugin\Context\ContextProviderInterface;
+use Drupal\page_manager\PageVariantInterface;
+
+/**
+ * Defines an interface for providing plugin contexts that require a page.
+ */
+interface PageVariantContextProviderInterface extends ContextProviderInterface {
+
+  /**
+   * {@inheritdoc}
+   *
+   * @param \Drupal\page_manager\PageVariantInterface $page_variant
+   *   The page variant to get contexts for.
+   */
+  public function getRuntimeContexts(array $unqualified_context_ids, PageVariantInterface $page_variant = NULL);
+
+  /**
+   * {@inheritdoc}
+   *
+   * @param \Drupal\page_manager\PageVariantInterface $page_variant
+   *   The page variant to get contexts for.
+   */
+  public function getAvailableContexts(PageVariantInterface $page_variant = NULL);
+
+}
diff --git a/src/ContextProvider/RouteParamContext.php b/src/ContextProvider/RouteParamContext.php
index b9a9952..d7b252f 100644
--- a/src/ContextProvider/RouteParamContext.php
+++ b/src/ContextProvider/RouteParamContext.php
@@ -12,13 +12,13 @@
 use Drupal\Core\StringTranslation\StringTranslationTrait;
 use Drupal\Core\Plugin\Context\Context;
 use Drupal\Core\Routing\RouteProviderInterface;
-use Drupal\page_manager\PageInterface;
+use Drupal\page_manager\PageVariantInterface;
 use Symfony\Component\HttpFoundation\RequestStack;
 
 /**
  * Sets values from the route parameters as a context.
  */
-class RouteParamContext implements PageContextProviderInterface {
+class RouteParamContext implements PageVariantContextProviderInterface {
 
   use StringTranslationTrait;
 
@@ -52,12 +52,13 @@ public function __construct(RouteProviderInterface $route_provider, RequestStack
   /**
    * {@inheritdoc}
    */
-  public function getRuntimeContexts(array $unqualified_context_ids, PageInterface $page = NULL) {
-    if (!$page) {
+  public function getRuntimeContexts(array $unqualified_context_ids, PageVariantInterface $page_variant = NULL) {
+    if (!$page_variant) {
       // @todo Add an assert here.
       return [];
     }
 
+    $page = $page_variant->getPage();
     $request = $this->requestStack->getCurrentRequest();
     $routes = $this->routeProvider->getRoutesByPattern($page->getPath())->all();
     $route = reset($routes);
@@ -95,8 +96,8 @@ public function getRuntimeContexts(array $unqualified_context_ids, PageInterface
   /**
    * {@inheritdoc}
    */
-  public function getAvailableContexts(PageInterface $page = NULL) {
-    return $this->getRuntimeContexts([], $page);
+  public function getAvailableContexts(PageVariantInterface $page_variant = NULL) {
+    return $this->getRuntimeContexts([], $page_variant);
   }
 
 }
diff --git a/src/ContextProvider/StaticContextProvider.php b/src/ContextProvider/StaticContextProvider.php
new file mode 100644
index 0000000..daa79fb
--- /dev/null
+++ b/src/ContextProvider/StaticContextProvider.php
@@ -0,0 +1,51 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\page_manager\ContextProvider\StaticContextProvider.
+ */
+
+namespace Drupal\page_manager\ContextProvider;
+
+use Drupal\page_manager\ContextMapperInterface;
+use Drupal\page_manager\PageVariantInterface;
+
+/**
+ * @todo.
+ */
+class StaticContextProvider implements PageVariantContextProviderInterface  {
+
+  /**
+   * @var \Drupal\page_manager\ContextMapperInterface
+   */
+  protected $contextMapper;
+
+  /**
+   * StaticContextProvider constructor.
+   *
+   * @param \Drupal\page_manager\ContextMapperInterface $context_mapper
+   */
+  public function __construct(ContextMapperInterface $context_mapper) {
+    $this->contextMapper = $context_mapper;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getRuntimeContexts(array $unqualified_context_ids, PageVariantInterface $page_variant = NULL) {
+    if (!$page_variant) {
+      // @todo Add an assert here.
+      return [];
+    }
+
+    return $this->contextMapper->getContextValues($page_variant->getStaticContexts());
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getAvailableContexts(PageVariantInterface $page_variant = NULL) {
+    return $this->getRuntimeContexts([], $page_variant);
+  }
+
+}
diff --git a/src/Entity/EntityWithContextAwarePluginCollectionTrait.php b/src/Entity/EntityWithContextAwarePluginCollectionTrait.php
deleted file mode 100644
index c7feef1..0000000
--- a/src/Entity/EntityWithContextAwarePluginCollectionTrait.php
+++ /dev/null
@@ -1,40 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\page_manager\Entity\EntityWithContextAwarePluginCollectionTrait.
- */
-
-namespace Drupal\page_manager\Entity;
-
-use Drupal\Core\Entity\EntityWithPluginCollectionInterface;
-use Drupal\Core\Plugin\ContextAwarePluginInterface;
-
-/**
- * Provides a trait for entities with context-aware plugins.
- */
-trait EntityWithContextAwarePluginCollectionTrait {
-
-  /**
-   * Finds the context mappings of all plugins contained by an entity.
-   *
-   * @return string[]
-   *   An array of unqualified context IDs.
-   *
-   * @see \Drupal\Core\Plugin\Context\ContextProviderInterface::getRuntimeContexts()
-   */
-  protected function getAllContextMappingIds() {
-    $context_mapping_ids = [];
-    if ($this instanceof EntityWithPluginCollectionInterface) {
-      foreach ($this->getPluginCollections() as $collection) {
-        foreach ($collection as $plugin) {
-          if ($plugin instanceof ContextAwarePluginInterface) {
-            $context_mapping_ids += array_values($plugin->getContextMapping());
-          }
-        }
-      }
-    }
-    return array_unique($context_mapping_ids);
-  }
-
-}
diff --git a/src/Entity/Page.php b/src/Entity/Page.php
index 6ae14fc..394f09d 100644
--- a/src/Entity/Page.php
+++ b/src/Entity/Page.php
@@ -7,6 +7,7 @@
 
 namespace Drupal\page_manager\Entity;
 
+use Drupal\page_manager\ObjectWithContextAwarePluginCollectionTrait;
 use Drupal\page_manager\PageInterface;
 use Drupal\Core\Condition\ConditionPluginCollection;
 use Drupal\Core\Config\Entity\ConfigEntityBase;
@@ -41,7 +42,7 @@
  */
 class Page extends ConfigEntityBase implements PageInterface {
 
-  use EntityWithContextAwarePluginCollectionTrait;
+  use ObjectWithContextAwarePluginCollectionTrait;
 
   /**
    * The ID of the page entity.
@@ -285,15 +286,15 @@ protected function filterParameters() {
    * {@inheritdoc}
    */
   public function getRuntimeContexts(array $context_ids = []) {
-    $context_ids = array_unique($context_ids + $this->getAllContextMappingIds());
-    return $this->contextRepository()->getRuntimeContexts($context_ids, $this);
+    $context_ids = array_unique($context_ids + $this->getAllContextMappingIds($this));
+    return $this->contextRepository()->getRuntimeContexts($context_ids);
   }
 
   /**
    * {@inheritdoc}
    */
   public function getAvailableContexts() {
-    return $this->contextRepository()->getAvailableContexts($this);
+    return $this->contextRepository()->getAvailableContexts();
   }
 
   /**
diff --git a/src/Entity/PageVariant.php b/src/Entity/PageVariant.php
index 97a4490..d9edca2 100644
--- a/src/Entity/PageVariant.php
+++ b/src/Entity/PageVariant.php
@@ -13,6 +13,7 @@
 use Drupal\Core\Entity\EntityStorageInterface;
 use Drupal\Core\Entity\EntityTypeInterface;
 use Drupal\Core\Plugin\DefaultSingleLazyPluginCollection;
+use Drupal\page_manager\ObjectWithContextAwarePluginCollectionTrait;
 use Drupal\page_manager\PageVariantInterface;
 
 /**
@@ -50,7 +51,7 @@
  */
 class PageVariant extends ConfigEntityBase implements PageVariantInterface {
 
-  use EntityWithContextAwarePluginCollectionTrait;
+  use ObjectWithContextAwarePluginCollectionTrait;
 
   /**
    * The ID of the page variant entity.
@@ -255,9 +256,8 @@ public function getPage() {
    */
   public function getRuntimeContexts() {
     if (is_null($this->runtimeContexts)) {
-      $static_contexts = $this->getContextMapper()->getContextValues($this->getStaticContexts());
-      $page_contexts = $this->getPage()->getRuntimeContexts($this->getAllContextMappingIds());
-      $this->runtimeContexts = array_merge($static_contexts, $page_contexts);
+      $context_mapping_ids = $this->getAllContextMappingIds($this) + $this->getAllContextMappingIds($this->getVariantPlugin());
+      $this->runtimeContexts = $this->contextRepository()->getRuntimeContexts($context_mapping_ids, $this);
     }
     return $this->runtimeContexts;
   }
@@ -267,9 +267,7 @@ public function getRuntimeContexts() {
    */
   public function getAvailableContexts() {
     if (is_null($this->availableContexts)) {
-      $static_contexts = $this->getContextMapper()->getContextValues($this->getStaticContexts());
-      $page_contexts = $this->getPage()->getAvailableContexts();
-      $this->availableContexts = array_merge($static_contexts, $page_contexts);
+      $this->availableContexts = $this->contextRepository()->getAvailableContexts($this);
     }
     return $this->availableContexts;
   }
@@ -437,6 +435,16 @@ protected function getContextMapper() {
   }
 
   /**
+   * Wraps the Page Manager context repository.
+   *
+   * @return \Drupal\page_manager\Context\PageManagerContextRepository
+   *   The Page Manager context repository.
+   */
+  protected function contextRepository() {
+    return \Drupal::service('page_manager.context_repository');
+  }
+
+  /**
    * Wraps the page entity storage.
    *
    * @return \Drupal\Core\Entity\EntityStorageInterface
diff --git a/src/ObjectWithContextAwarePluginCollectionTrait.php b/src/ObjectWithContextAwarePluginCollectionTrait.php
new file mode 100644
index 0000000..9fd4e6d
--- /dev/null
+++ b/src/ObjectWithContextAwarePluginCollectionTrait.php
@@ -0,0 +1,45 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\page_manager\ObjectWithContextAwarePluginCollectionTrait.
+ */
+
+namespace Drupal\page_manager;
+
+use Drupal\Core\Entity\EntityWithPluginCollectionInterface;
+use Drupal\Core\Plugin\ContextAwarePluginInterface;
+
+/**
+ * Provides a trait for entities with context-aware plugins.
+ *
+ * @todo Move to CTools.
+ */
+trait ObjectWithContextAwarePluginCollectionTrait {
+
+  /**
+   * Finds the context mappings of all plugins contained by an entity.
+   *
+   * @param object $object
+   *   An object to inspect for plugin collections.
+   *
+   * @return string[]
+   *   An array of unqualified context IDs.
+   *
+   * @see \Drupal\Core\Plugin\Context\ContextProviderInterface::getRuntimeContexts()
+   */
+  protected function getAllContextMappingIds($object) {
+    $context_mapping_ids = [];
+    if ($object instanceof EntityWithPluginCollectionInterface || $object instanceof ObjectWithPluginCollectionInterface) {
+      foreach ($object->getPluginCollections() as $key => $collection) {
+        foreach ($collection as $plugin) {
+          if ($plugin instanceof ContextAwarePluginInterface) {
+            $context_mapping_ids = array_merge($context_mapping_ids, array_values($plugin->getContextMapping()));
+          }
+        }
+      }
+    }
+    return array_values(array_unique($context_mapping_ids));
+  }
+
+}
diff --git a/src/ObjectWithPluginCollectionInterface.php b/src/ObjectWithPluginCollectionInterface.php
new file mode 100644
index 0000000..98a1983
--- /dev/null
+++ b/src/ObjectWithPluginCollectionInterface.php
@@ -0,0 +1,28 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\page_manager\ObjectWithPluginCollectionInterface.
+ */
+
+namespace Drupal\page_manager;
+
+/**
+ * @todo.
+ *
+ * @see \Drupal\Core\Entity\EntityWithPluginCollectionInterface
+ *
+ * @todo Move to CTools or core.
+ */
+interface ObjectWithPluginCollectionInterface {
+
+  /**
+   * Gets the plugin collections used by this entity.
+   *
+   * @return \Drupal\Component\Plugin\LazyPluginCollection[]
+   *   An array of plugin collections, keyed by the property name they use to
+   *   store their configuration.
+   */
+  public function getPluginCollections();
+
+}
diff --git a/src/Plugin/DisplayVariant/PageBlockDisplayVariant.php b/src/Plugin/DisplayVariant/PageBlockDisplayVariant.php
index 78f6dc6..0e94c2a 100644
--- a/src/Plugin/DisplayVariant/PageBlockDisplayVariant.php
+++ b/src/Plugin/DisplayVariant/PageBlockDisplayVariant.php
@@ -21,6 +21,7 @@
 use Drupal\Core\Session\AccountInterface;
 use Drupal\Core\Utility\Token;
 use Drupal\ctools\Plugin\DisplayVariant\BlockDisplayVariant;
+use Drupal\page_manager\ObjectWithPluginCollectionInterface;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
@@ -31,7 +32,7 @@
  *   admin_label = @Translation("Block page")
  * )
  */
-class PageBlockDisplayVariant extends BlockDisplayVariant {
+class PageBlockDisplayVariant extends BlockDisplayVariant implements ObjectWithPluginCollectionInterface {
 
   /**
    * The module handler.
@@ -300,4 +301,13 @@ public function getRegionNames() {
     ];
   }
 
+  /**
+   * {@inheritdoc}
+   */
+  public function getPluginCollections() {
+    return [
+      'blocks' => $this->getBlockCollection(),
+    ];
+  }
+
 }
diff --git a/src/Routing/VariantRouteFilter.php b/src/Routing/VariantRouteFilter.php
index f588abd..6d2abcc 100644
--- a/src/Routing/VariantRouteFilter.php
+++ b/src/Routing/VariantRouteFilter.php
@@ -153,7 +153,7 @@ protected function routeWeightSort(Route $a, Route $b) {
    *   TRUE if the route is valid, FALSE otherwise.
    */
   protected function checkPageVariantAccess($page_variant_id) {
-    /** @var \Drupal\page_manager\PageVariantInterface $page */
+    /** @var \Drupal\page_manager\PageVariantInterface $variant */
     $variant = $this->pageVariantStorage->load($page_variant_id);
 
     try {
diff --git a/src/Tests/StaticContextTest.php b/src/Tests/StaticContextTest.php
index 38ba04f..8aae7fd 100644
--- a/src/Tests/StaticContextTest.php
+++ b/src/Tests/StaticContextTest.php
@@ -110,7 +110,7 @@ public function testStaticContext() {
       'settings[label_display]' => 1,
       'settings[view_mode]' => 'default',
       'region' => 'bottom',
-      'context_mapping[entity]' => $contexts[1]['machine_name'],
+      'context_mapping[entity]' => '@page_manager.static_context_provider:static_node_2',
     ];
     $this->drupalPostForm(NULL, $edit, 'Add block');
     $this->assertText($edit['settings[label]']);
diff --git a/tests/src/Unit/ObjectWithContextAwarePluginCollectionTraitTest.php b/tests/src/Unit/ObjectWithContextAwarePluginCollectionTraitTest.php
new file mode 100644
index 0000000..50c5f52
--- /dev/null
+++ b/tests/src/Unit/ObjectWithContextAwarePluginCollectionTraitTest.php
@@ -0,0 +1,104 @@
+<?php
+
+/**
+ * @file
+ * Contains
+ *   \Drupal\Tests\page_manager\Unit\ObjectWithContextAwarePluginCollectionTraitTest.
+ */
+
+namespace Drupal\Tests\page_manager\Unit;
+
+use Drupal\Component\Plugin\PluginManagerInterface;
+use Drupal\Core\Plugin\ContextAwarePluginInterface;
+use Drupal\Core\Plugin\DefaultLazyPluginCollection;
+use Drupal\page_manager\ObjectWithContextAwarePluginCollectionTrait;
+use Drupal\page_manager\ObjectWithPluginCollectionInterface;
+use Drupal\Tests\Core\Config\Entity\Fixtures\ConfigEntityBaseWithPluginCollections;
+use Drupal\Tests\UnitTestCase;
+
+/**
+ * @coversDefaultClass \Drupal\page_manager\ObjectWithContextAwarePluginCollectionTrait
+ *
+ * @group PageManager
+ */
+class ObjectWithContextAwarePluginCollectionTraitTest extends UnitTestCase{
+
+  /**
+   * @covers ::getAllContextMappingIds
+   * @dataProvider providerTestGetAllContextMappingIds
+   */
+  public function testGetAllContextMappingIds($object) {
+    $plugin1 = $this->prophesize(ContextAwarePluginInterface::class);
+    $plugin1->getContextMapping()->willReturn(['foo' => 'John']);
+    $plugin2 = $this->prophesize(ContextAwarePluginInterface::class);
+    $plugin2->getContextMapping()->willReturn(['foo' => 'John', 'bar' => 'Paul']);
+    $plugin3 = $this->prophesize(ContextAwarePluginInterface::class);
+    $plugin3->getContextMapping()->willReturn(['foo' => 'George', 'bar' => 'Ringo']);
+    $plugin4 = $this->prophesize(ContextAwarePluginInterface::class);
+    $plugin4->getContextMapping()->willReturn(['foo' => 'Ringo', 'bar' => 'George']);
+
+    $plugin_manager = $this->prophesize(PluginManagerInterface::class);
+    $plugin_manager->createInstance('plugin_id1', ['id' => 'plugin_id1'])->willReturn($plugin1->reveal());
+    $plugin_manager->createInstance('plugin_id2', ['id' => 'plugin_id2'])->willReturn($plugin2->reveal());
+    $plugin_manager->createInstance('plugin_id3', ['id' => 'plugin_id3'])->willReturn($plugin3->reveal());
+    $plugin_manager->createInstance('plugin_id4', ['id' => 'plugin_id4'])->willReturn($plugin4->reveal());
+
+    $object->setPluginManager($plugin_manager->reveal());
+
+    $expected = ['John', 'Paul', 'George', 'Ringo'];
+    $context_mapping_ids = $object->testGetAllContextMappingIds();
+    $this->assertSame($expected, $context_mapping_ids);
+  }
+
+  public function providerTestGetAllContextMappingIds() {
+    $data = [];
+    $data['entity'] = [new TestEntity([], '')];
+    $data['object'] = [new TestObject()];
+    return $data;
+  }
+
+}
+
+class TestObject implements ObjectWithPluginCollectionInterface {
+
+  use ObjectWithContextAwarePluginCollectionTrait;
+  use HelperTrait;
+
+}
+
+class TestEntity extends ConfigEntityBaseWithPluginCollections {
+
+  use ObjectWithContextAwarePluginCollectionTrait;
+  use HelperTrait;
+
+}
+
+trait HelperTrait {
+
+  protected $pluginManager;
+
+  public function testGetAllContextMappingIds() {
+    return $this->getAllContextMappingIds($this);
+  }
+
+  public function setPluginManager(PluginManagerInterface $plugin_manager) {
+    $this->pluginManager = $plugin_manager;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getPluginCollections() {
+    return [
+      'plugin_collection1' => new DefaultLazyPluginCollection($this->pluginManager, [
+        'plugin_id1' => ['id' => 'plugin_id1'],
+        'plugin_id2' => ['id' => 'plugin_id2'],
+      ]),
+      'plugin_collection2' => new DefaultLazyPluginCollection($this->pluginManager, [
+        'plugin_id3' => ['id' => 'plugin_id3'],
+        'plugin_id4' => ['id' => 'plugin_id4'],
+      ]),
+    ];
+  }
+
+}
diff --git a/tests/src/Unit/PageTest.php b/tests/src/Unit/PageTest.php
index 2017218..00d58c8 100644
--- a/tests/src/Unit/PageTest.php
+++ b/tests/src/Unit/PageTest.php
@@ -79,7 +79,7 @@ public function testGetRuntimeContexts() {
     $context = new Context(new ContextDefinition('bar'));
 
     $context_repository = $this->prophesize(ContextRepositoryInterface::class);
-    $context_repository->getRuntimeContexts([], $this->page)->willReturn(['foo' => $context]);
+    $context_repository->getRuntimeContexts([])->willReturn(['foo' => $context]);
 
     $condition_plugin_manager = $this->prophesize(PluginManagerInterface::class);
 
diff --git a/tests/src/Unit/PageVariantTest.php b/tests/src/Unit/PageVariantTest.php
index d0e71f4..4274f96 100644
--- a/tests/src/Unit/PageVariantTest.php
+++ b/tests/src/Unit/PageVariantTest.php
@@ -11,6 +11,7 @@
 use Drupal\Core\DependencyInjection\ContainerBuilder;
 use Drupal\Core\Entity\EntityStorageInterface;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
+use Drupal\Core\Plugin\Context\ContextRepositoryInterface;
 use Drupal\page_manager\ContextMapperInterface;
 use Drupal\page_manager\Entity\PageVariant;
 use Drupal\page_manager\PageInterface;
@@ -34,9 +35,9 @@ class PageVariantTest extends UnitTestCase {
   protected $page;
 
   /**
-   * @var \Drupal\page_manager\ContextMapperInterface|\Prophecy\Prophecy\ProphecyInterface
+   * @var \Drupal\Core\Plugin\Context\ContextRepositoryInterface|\Prophecy\Prophecy\ProphecyInterface
    */
-  protected $contextMapper;
+  protected $contextRepository;
 
   /**
    * {@inheritdoc}
@@ -45,24 +46,16 @@ protected function setUp() {
     parent::setUp();
 
     $this->pageVariant = new PageVariant(['id' => 'the_page_variant', 'page' => 'the_page'], 'page_variant');
-    $this->page = $this->prophesize(PageInterface::class);
 
-    $entity_storage = $this->prophesize(EntityStorageInterface::class);
-    $entity_storage->load('the_page')->willReturn($this->page->reveal());
-
-    $entity_type_manager = $this->prophesize(EntityTypeManagerInterface::class);
-    $entity_type_manager->getStorage('page')->willReturn($entity_storage);
-
-    $this->contextMapper = $this->prophesize(ContextMapperInterface::class);
+    $this->contextRepository = $this->prophesize(ContextRepositoryInterface::class);
 
     $condition_plugin_manager = $this->prophesize(PluginManagerInterface::class);
     $display_variant_plugin_manager = $this->prophesize(PluginManagerInterface::class);
 
     $container = new ContainerBuilder();
-    $container->set('entity_type.manager', $entity_type_manager->reveal());
-    $container->set('page_manager.context_mapper', $this->contextMapper->reveal());
     $container->set('plugin.manager.condition', $condition_plugin_manager->reveal());
     $container->set('plugin.manager.display_variant', $display_variant_plugin_manager->reveal());
+    $container->set('page_manager.context_repository', $this->contextRepository->reveal());
     \Drupal::setContainer($container);
   }
 
@@ -71,8 +64,10 @@ protected function setUp() {
    * @dataProvider providerTestGetRuntimeContexts
    */
   public function testGetRuntimeContexts($static_contexts, $page_contexts, $expected) {
-    $this->contextMapper->getContextValues([])->willReturn($static_contexts)->shouldBeCalledTimes(1);
-    $this->page->getRuntimeContexts([])->willReturn($page_contexts)->shouldBeCalledTimes(1);
+    $this->contextRepository
+      ->getRuntimeContexts([], $this->pageVariant)
+      ->willReturn(array_merge($static_contexts, $page_contexts))
+      ->shouldBeCalledTimes(1);
 
     $contexts = $this->pageVariant->getRuntimeContexts();
     $this->assertSame($expected, $contexts);
@@ -105,8 +100,7 @@ public function providerTestGetRuntimeContexts() {
    * @covers ::removeStaticContext
    */
   public function testGetRuntimeContextsAfterReset() {
-    $this->contextMapper->getContextValues([])->willReturn([])->shouldBeCalledTimes(2);
-    $this->page->getRuntimeContexts([])->willReturn([])->shouldBeCalledTimes(2);
+    $this->contextRepository->getRuntimeContexts([], $this->pageVariant)->willReturn([])->shouldBeCalledTimes(2);
 
     $expected = [];
     $contexts = $this->pageVariant->getRuntimeContexts();
diff --git a/tests/src/Unit/RouteParamContextTest.php b/tests/src/Unit/RouteParamContextTest.php
index ab91db5..ef5334b 100644
--- a/tests/src/Unit/RouteParamContextTest.php
+++ b/tests/src/Unit/RouteParamContextTest.php
@@ -14,6 +14,7 @@
 use Drupal\Core\TypedData\TypedDataManager;
 use Drupal\page_manager\ContextProvider\RouteParamContext;
 use Drupal\page_manager\PageInterface;
+use Drupal\page_manager\PageVariantInterface;
 use Drupal\Tests\UnitTestCase;
 use Prophecy\Argument;
 use Symfony\Component\HttpFoundation\Request;
@@ -44,6 +45,8 @@ class RouteParamContextTest extends UnitTestCase {
    */
   protected $page;
 
+  protected $pageVariant;
+
   /**
    * {@inheritdoc}
    */
@@ -58,6 +61,8 @@ protected function setUp() {
     \Drupal::setContainer($container);
 
     $this->page = $this->prophesize(PageInterface::class);
+    $this->pageVariant = $this->prophesize(PageVariantInterface::class);
+    $this->pageVariant->getPage()->willReturn($this->page->reveal());
   }
 
   /**
@@ -103,7 +108,7 @@ public function testGetRuntimeContexts() {
     $request->attributes->add(['foo' => 'banana']);
 
     $route_param_context = new RouteParamContext($route_provider->reveal(), $request_stack);
-    $contexts = $route_param_context->getRuntimeContexts([], $this->page->reveal());
+    $contexts = $route_param_context->getRuntimeContexts([], $this->pageVariant->reveal());
     $this->assertSame(['foo', 'baz', 'page'], array_keys($contexts));
     $this->assertSame('Foo', (string) $contexts['foo']->getContextDefinition()->getLabel());
     $this->assertSame('{baz} from route', (string) $contexts['baz']->getContextDefinition()->getLabel());
@@ -132,7 +137,7 @@ public function testGetRuntimeContextsEmpty() {
     $request->attributes->add(['foo' => 'banana']);
 
     $route_param_context = new RouteParamContext($route_provider->reveal(), $request_stack);
-    $contexts = $route_param_context->getRuntimeContexts([], $this->page->reveal());
+    $contexts = $route_param_context->getRuntimeContexts([], $this->pageVariant->reveal());
     $this->assertEmpty($contexts);
   }