diff --git a/src/PanelsStorage.php b/src/PanelsStorage.php
index 24cb76a..48882e0 100644
--- a/src/PanelsStorage.php
+++ b/src/PanelsStorage.php
@@ -73,7 +73,7 @@ public function load($id) {
       // If this page variant doesn't have a Panels display on it, then we treat
       // it the same as if there was no such page variant.
       if (!($panels_display instanceof PanelsDisplayVariant)) {
-        return;
+        return NULL;
       }
 
       // Pass down the contexts because the display has no other way to get them
@@ -96,4 +96,3 @@ public function access($id, $op, AccountInterface $account) {
   }
 
 }
-
diff --git a/tests/src/Unit/PanelsStorageTest.php b/tests/src/Unit/PanelsStorageTest.php
index cec2e06..a1047a8 100644
--- a/tests/src/Unit/PanelsStorageTest.php
+++ b/tests/src/Unit/PanelsStorageTest.php
@@ -8,7 +8,13 @@
 namespace Drupal\Tests\page_manager\Unit;
 
 use Drupal\Core\Access\AccessResult;
+use Drupal\Core\Entity\EntityStorageInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
+use Drupal\Core\Session\AccountInterface;
+use Drupal\page_manager\PageVariantInterface;
 use Drupal\page_manager\PanelsStorage;
+use Drupal\page_manager\Plugin\DisplayVariant\HttpStatusCodeDisplayVariant;
+use Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant;
 use Drupal\Tests\UnitTestCase;
 
 /**
@@ -21,32 +27,27 @@
 class PanelsStorageTest extends UnitTestCase {
 
   /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant
+   * @var \Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant|\Prophecy\Prophecy\ProphecyInterface
    */
   protected $panelsDisplay;
 
   /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\page_manager\Plugin\DisplayVariant\HttpStatusCodeDisplayVariant
-   */
-  protected $httpStatusCode;
-
-  /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\page_manager\PageVariantInterface $page_variant
+   * @var \Drupal\page_manager\PageVariantInterface|\Prophecy\Prophecy\ProphecyInterface
    */
   protected $pageVariant;
 
   /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\page_manager\PageVariantInterface
+   * @var \Drupal\page_manager\PageVariantInterface|\Prophecy\Prophecy\ProphecyInterface
    */
   protected $pageVariantNotPanels;
 
   /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\Core\Entity\EntityStorageInterface
+   * @var \Drupal\Core\Entity\EntityStorageInterface|\Prophecy\Prophecy\ProphecyInterface
    */
   protected $storage;
 
   /**
-   * @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\Core\Entity\EntityTypeManagerInterface
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface|\Prophecy\Prophecy\ProphecyInterface
    */
   protected $entityTypeManager;
 
@@ -56,31 +57,21 @@ class PanelsStorageTest extends UnitTestCase {
   protected function setUp() {
     parent::setUp();
 
-    $this->panelsDisplay = $this->getMockBuilder('Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant')
-      ->disableOriginalConstructor()
-      ->getMock();
-
-    $this->httpStatusCode = $this->getMockBuilder('Drupal\page_manager\Plugin\DisplayVariant\HttpStatusCodeDisplayVariant')
-      ->disableOriginalConstructor()
-      ->getMock();
+    $this->panelsDisplay = $this->prophesize(PanelsDisplayVariant::class);
 
-    $this->pageVariant = $this->getMock('Drupal\page_manager\PageVariantInterface');
-    $this->pageVariant->method('getVariantPlugin')->willReturn($this->panelsDisplay);
+    $this->pageVariant = $this->prophesize(PageVariantInterface::class);
+    $this->pageVariant->getVariantPlugin()->willReturn($this->panelsDisplay->reveal());
 
-    $this->pageVariantNotPanels = $this->getMock('Drupal\page_manager\PageVariantInterface');
-    $this->pageVariantNotPanels->method('getVariantPlugin')->willReturn($this->httpStatusCode);
-    $this->pageVariantNotPanels->expects($this->never())->method('getContexts');
+    $this->pageVariantNotPanels = $this->prophesize(PageVariantInterface::class);
+    $this->pageVariantNotPanels->getContexts()->shouldNotBeCalled();
 
-    $this->storage = $this->getMock('Drupal\Core\Entity\EntityStorageInterface');
-    $this->storage->method('load')->willReturnMap([
-      ['id_exists',    $this->pageVariant],
-      ['doesnt_exist', NULL],
-      ['not_a_panel',  $this->pageVariantNotPanels],
-    ]);
+    $non_panels_variant = $this->prophesize(HttpStatusCodeDisplayVariant::class);
+    $this->pageVariantNotPanels->getVariantPlugin()->willReturn($non_panels_variant->reveal());
 
-    $this->entityTypeManager = $this->getMock('Drupal\Core\Entity\EntityTypeManagerInterface');
-    $this->entityTypeManager->method('getStorage')->with('page_variant')->willReturn($this->storage);
+    $this->storage = $this->prophesize(EntityStorageInterface::class);
 
+    $this->entityTypeManager = $this->prophesize(EntityTypeManagerInterface::class);
+    $this->entityTypeManager->getStorage('page_variant')->willReturn($this->storage->reveal());
   }
 
   /**
@@ -88,20 +79,23 @@ protected function setUp() {
    */
   public function testLoad() {
     // Make sure that the contexts are passed down (or not).
-    $this->pageVariant->expects($this->once())->method('getContexts')->willReturn([]);
-    $this->panelsDisplay->expects($this->once())->method('setContexts')->with([]);
-    $this->pageVariantNotPanels->expects($this->never())->method('getContexts');
+    $this->pageVariant->getContexts()->willReturn([]);
+    $this->panelsDisplay->setContexts([])->shouldBeCalledTimes(1);
 
-    $panels_storage = new PanelsStorage($this->entityTypeManager);
+    $this->storage->load('id_exists')->willReturn($this->pageVariant->reveal());
+    $this->storage->load('doesnt_exist')->willReturn(NULL);
+    $this->storage->load('not_a_panel')->willReturn($this->pageVariantNotPanels->reveal());
+
+    $panels_storage = new PanelsStorage($this->entityTypeManager->reveal());
 
     // Test the success condition.
-    $this->assertEquals($this->panelsDisplay, $panels_storage->load('id_exists'));
+    $this->assertSame($this->panelsDisplay->reveal(), $panels_storage->load('id_exists'));
 
     // Should be NULL if it doesn't exist.
-    $this->assertEquals(NULL, $panels_storage->load('doesnt_exist'));
+    $this->assertNull($panels_storage->load('doesnt_exist'));
 
     // Should also be NULL if it's not a PanelsDisplayVariant.
-    $this->assertEquals(NULL, $panels_storage->load('not_a_panel'));
+    $this->assertNull($panels_storage->load('not_a_panel'));
   }
 
   /**
@@ -109,18 +103,17 @@ public function testLoad() {
    */
   public function testSaveSuccessful() {
     $test_config = ['my_config' => '123'];
-    $this->panelsDisplay->expects($this->once())->method('setConfiguration')->with($test_config);
-    $this->pageVariant->expects($this->once())->method('save');
-
-    /** @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant $panels_display */
-    $panels_display = $this->getMockBuilder('Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant')
-      ->disableOriginalConstructor()
-      ->getMock();
-    $panels_display->expects($this->once())->method('getStorageId')->willReturn('id_exists');
-    $panels_display->expects($this->once())->method('getConfiguration')->willReturn($test_config);
-
-    $panels_storage = new PanelsStorage($this->entityTypeManager);
-    $panels_storage->save($panels_display);
+    $this->panelsDisplay->setConfiguration($test_config)->shouldBeCalledTimes(1);
+    $this->pageVariant->save()->shouldBeCalledTimes(1);
+
+    $this->storage->load('id_exists')->willReturn($this->pageVariant->reveal());
+
+    $panels_display = $this->prophesize(PanelsDisplayVariant::class);
+    $panels_display->getStorageId()->willReturn('id_exists');
+    $panels_display->getConfiguration()->willReturn($test_config);
+
+    $panels_storage = new PanelsStorage($this->entityTypeManager->reveal());
+    $panels_storage->save($panels_display->reveal());
   }
 
   /**
@@ -130,18 +123,17 @@ public function testSaveSuccessful() {
    * @expectedExceptionMessage Couldn't find page variant to store Panels display
    */
   public function testSaveDoesntExist() {
-    $this->panelsDisplay->expects($this->never())->method('setConfiguration');
-    $this->pageVariant->expects($this->never())->method('save');
-
-    /** @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant $panels_display */
-    $panels_display = $this->getMockBuilder('Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant')
-      ->disableOriginalConstructor()
-      ->getMock();
-    $panels_display->expects($this->once())->method('getStorageId')->willReturn('doesnt_exist');
-    $panels_display->expects($this->never())->method('getConfiguration');
-
-    $panels_storage = new PanelsStorage($this->entityTypeManager);
-    $panels_storage->save($panels_display);
+    $this->panelsDisplay->setConfiguration()->shouldNotBeCalled();
+    $this->pageVariant->save()->shouldNotBeCalled();
+
+    $this->storage->load('doesnt_exist')->willReturn(NULL);
+
+    $panels_display = $this->prophesize(PanelsDisplayVariant::class);
+    $panels_display->getStorageId()->willReturn('doesnt_exist');
+    $panels_display->getConfiguration()->shouldNotBeCalled();
+
+    $panels_storage = new PanelsStorage($this->entityTypeManager->reveal());
+    $panels_storage->save($panels_display->reveal());
   }
 
   /**
@@ -151,39 +143,37 @@ public function testSaveDoesntExist() {
    * @expectedExceptionMessage Page variant doesn't use a Panels display variant
    */
   public function testSaveNotPanels() {
-    $this->panelsDisplay->expects($this->never())->method('setConfiguration');
-    $this->pageVariant->expects($this->never())->method('save');
-
-    /** @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant $panels_display */
-    $panels_display = $this->getMockBuilder('Drupal\panels\Plugin\DisplayVariant\PanelsDisplayVariant')
-      ->disableOriginalConstructor()
-      ->getMock();
-    $panels_display->expects($this->once())->method('getStorageId')->willReturn('not_a_panel');
-    $panels_display->expects($this->never())->method('getConfiguration');
-
-    $panels_storage = new PanelsStorage($this->entityTypeManager);
-    $panels_storage->save($panels_display);
+    $this->storage->load('not_a_panel')->willReturn($this->pageVariantNotPanels->reveal());
+
+    $this->panelsDisplay->setConfiguration()->shouldNotBeCalled();
+    $this->pageVariant->save()->shouldNotBeCalled();
+
+    $panels_display = $this->prophesize(PanelsDisplayVariant::class);
+    $panels_display->getStorageId()->willReturn('not_a_panel');
+    $panels_display->getConfiguration()->shouldNotBeCalled();
+
+    $panels_storage = new PanelsStorage($this->entityTypeManager->reveal());
+    $panels_storage->save($panels_display->reveal());
   }
 
   /**
    * @covers ::access
    */
   public function testAccess() {
-    /** @var \PHPUnit_Framework_MockObject_MockObject|\Drupal\Core\Session\AccountInterface $account */
-    $account = $this->getMock('\Drupal\Core\Session\AccountInterface');
+    $this->storage->load('id_exists')->willReturn($this->pageVariant->reveal());
+    $this->storage->load('doesnt_exist')->willReturn(NULL);
+
+    $account = $this->prophesize(AccountInterface::class);
 
-    $this->pageVariant->expects($this->once())
-      ->method('access')
-      ->with('read', $account, TRUE)
-      ->willReturn(AccessResult::allowed());
+    $this->pageVariant->access('read', $account->reveal(), TRUE)->willReturn(AccessResult::allowed());
 
-    $panels_storage = new PanelsStorage($this->entityTypeManager);
+    $panels_storage = new PanelsStorage($this->entityTypeManager->reveal());
 
     // Test the access condition.
-    $this->assertEquals(AccessResult::allowed(), $panels_storage->access('id_exists', 'read', $account));
+    $this->assertEquals(AccessResult::allowed(), $panels_storage->access('id_exists', 'read', $account->reveal()));
 
     // Should be forbidden if it doesn't exist.
-    $this->assertEquals(AccessResult::forbidden(), $panels_storage->access('doesnt_exist', 'read', $account));
+    $this->assertEquals(AccessResult::forbidden(), $panels_storage->access('doesnt_exist', 'read', $account->reveal()));
   }
 
 }