diff --git a/payment/lib/Drupal/payment/Generate.php b/payment/lib/Drupal/payment/Generate.php
index 19554c3..68601d4 100644
--- a/payment/lib/Drupal/payment/Generate.php
+++ b/payment/lib/Drupal/payment/Generate.php
@@ -27,10 +27,10 @@ class Generate {
    */
   static function createPayment($uid, PaymentMethod $payment_method = NULL) {
     $context_manager = \Drupal::service('plugin.manager.payment.context');
-    $payment = entity_create('payment', array())
-      ->setPaymentMethodId('payment_unavailable')
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ))->setPaymentMethodId('payment_unavailable')
       ->setOwnerId($uid)
-      ->setPaymentContext($context_manager->createInstance('payment_unavailable'))
       ->setLineItems(static::createPaymentLineItems());
 
     return $payment;
diff --git a/payment/lib/Drupal/payment/Plugin/Core/Entity/Payment.php b/payment/lib/Drupal/payment/Plugin/Core/Entity/Payment.php
index 42daba4..e9f121b 100644
--- a/payment/lib/Drupal/payment/Plugin/Core/Entity/Payment.php
+++ b/payment/lib/Drupal/payment/Plugin/Core/Entity/Payment.php
@@ -26,6 +26,7 @@ use Drupal\payment\Plugin\payment\status\PaymentStatusInterface;
  *     "storage" = "Drupal\payment\Plugin\Core\Entity\PaymentStorageController",
  *   },
  *   entity_keys = {
+ *     "bundle" = "bundle",
  *     "id" = "id",
  *     "uuid" = "uuid",
  *   },
@@ -63,6 +64,14 @@ class Payment extends EntityNG implements PaymentInterface {
   protected $statuses = array();
 
   /**
+   * Overrides Entity::__construct().
+   */
+  public function __construct(array $values, $entity_type, $bundle = FALSE, $translations = array()) {
+    parent::__construct($values, $entity_type, $bundle, $translations);
+    $this->context = \Drupal::service('plugin.manager.payment.context')->createInstance($this->bundle());
+  }
+
+  /**
    * {@inheritdoc}
    */
   public function label($langcode = NULL) {
@@ -76,15 +85,6 @@ class Payment extends EntityNG implements PaymentInterface {
   /**
    * {@inheritdoc}
    */
-  public function setPaymentContext(PaymentContextInterface $context) {
-    $this->context = $context;
-
-    return $this;
-  }
-
-  /**
-   * {@inheritdoc}
-   */
   public function getPaymentContext() {
     return $this->context;
   }
diff --git a/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentInterface.php b/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentInterface.php
index 677998d..6668dee 100644
--- a/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentInterface.php
+++ b/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentInterface.php
@@ -20,16 +20,6 @@ use Drupal\payment\Plugin\payment\status\PaymentStatusInterface;
 interface PaymentInterface extends EntityInterface, ExecutableInterface {
 
   /**
-   * Sets the machine name of the context that created this Payment, such as a
-   * payment form, or a module.
-   *
-   * @param \Drupal\payment\Plugin\payment\context\PaymentContextInterface $context
-   *
-   * @return \Drupal\payment\Plugin\Core\Entity\PaymentInterface
-   */
-  public function setPaymentContext(PaymentContextInterface $context);
-
-  /**
    * Gets the machine name of the context that created this Payment, such as a
    * payment form, or a module.
    *
diff --git a/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentStorageController.php b/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentStorageController.php
index d807ba7..a774b3e 100644
--- a/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentStorageController.php
+++ b/payment/lib/Drupal/payment/Plugin/Core/Entity/PaymentStorageController.php
@@ -19,6 +19,9 @@ class PaymentStorageController extends DatabaseStorageControllerNG implements Pa
    * {@inheritdoc}
    */
   function create(array $values) {
+    if (isset($values['context']) && !isset($values['bundle'])) {
+      $values['bundle'] = $values['context']->getPluginId();
+    }
     $payment = parent::create($values);
     $payment->setStatus(\Drupal::service('plugin.manager.payment.status')->createInstance('payment_created'));
 
@@ -34,7 +37,7 @@ class PaymentStorageController extends DatabaseStorageControllerNG implements Pa
     $statuses = $this->loadPaymentStatuses(array_keys($queried_entities));
     foreach ($queried_entities as $id => $queried_entity) {
       $queried_entities[$id] = (object) array(
-        'context' => $queried_entity->context_plugin_id ? $manager->createInstance($queried_entity->context_plugin_id) : NULL,
+        'bundle' => $queried_entity->bundle,
         'currencyCode' => $queried_entity->currency_code,
         'id' => (int) $queried_entity->id,
         'lineItems' => $line_items[$id],
@@ -91,7 +94,7 @@ class PaymentStorageController extends DatabaseStorageControllerNG implements Pa
    */
   protected function mapToStorageRecord(EntityInterface $entity) {
     $record = new \stdClass();
-    $record->context_plugin_id = $entity->getPaymentContext() ? $entity->getPaymentContext()->getPluginId() : NULL;
+    $record->bundle = $entity->bundle();
     $record->currency_code = $entity->id();
     $record->id = $entity->id();
     $record->payment_method_id = $entity->getPaymentMethodId();
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentAccessControllerUnitTest.php b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentAccessControllerUnitTest.php
index 41e7252..ac1e44f 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentAccessControllerUnitTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentAccessControllerUnitTest.php
@@ -41,7 +41,9 @@ class PaymentAccessControllerUnitTest extends AccessibleInterfaceUnitTestBase {
     ));
 
     // Create a new payment.
-    $this->assertDataAccess(entity_create('payment', array()), 'a payment', 'create', $authenticated, array(), array(
+    $this->assertDataAccess(entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    )), 'a payment', 'create', $authenticated, array(), array(
       'anonymous' => TRUE,
       'authenticated_without_permissions' => TRUE,
     ));
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentStorageControllerWebTest.php b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentStorageControllerWebTest.php
index 83c2ce3..0875587 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentStorageControllerWebTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentStorageControllerWebTest.php
@@ -35,15 +35,16 @@ class PaymentStorageControllerWebTest extends WebTestBase {
    */
   function testCRUD() {
     // Test creating a payment.
-    $payment = entity_create('payment', array());
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ));
     $this->assertTrue($payment instanceof PaymentInterface);
     $this->assertTrue(is_int($payment->getOwnerId()));
     $this->assertEqual(count($payment->validate()), 0);
+    $this->assertTrue($payment->getPaymentContext() instanceof PaymentContextInterface);
 
     // Test saving a payment.
-    $context_manager = $this->container->get('plugin.manager.payment.context');
     $this->assertFalse($payment->id());
-    $payment->setPaymentContext($context_manager->createInstance('payment_unavailable'));
     $payment->save();
     // @todo The ID should be an integer, but for some reason the entity field
     //   API returns a string.
@@ -53,7 +54,6 @@ class PaymentStorageControllerWebTest extends WebTestBase {
     $payment_loaded = entity_load_unchanged('payment', $payment->id());
     $this->assertEqual(count($payment_loaded->getLineItems()), count($payment->getLineItems()));
     $this->assertEqual(count($payment_loaded->getStatuses()), count($payment->getStatuses()));
-    $this->assertTrue($payment_loaded->getPaymentContext() instanceof PaymentContextInterface);
 
     // Test deleting a payment.
     $payment->delete();
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentUnitTest.php b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentUnitTest.php
index b8dc472..c996af1 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentUnitTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/Core/Entity/PaymentUnitTest.php
@@ -9,6 +9,7 @@ namespace Drupal\payment\Tests\Plugin\Core\Entity;
 
 use Drupal\payment\Plugin\Core\Entity\PaymentInterface;
 use Drupal\payment\Plugin\Core\Entity\PaymentMethodInterface;
+use Drupal\payment\Plugin\payment\Context\PaymentContextInterface;
 use Drupal\payment\Generate;
 use Drupal\simpletest\DrupalUnitTestBase;
 
@@ -35,7 +36,10 @@ class PaymentUnitTest extends DrupalUnitTestBase {
    */
   function setUp() {
     parent::setUp();
-    $this->payment = entity_create('payment', array());
+    $this->bundle = 'payment_unavailable';
+    $this->payment = entity_create('payment', array(
+      'bundle' => $this->bundle,
+    ));
   }
 
   /**
@@ -46,12 +50,19 @@ class PaymentUnitTest extends DrupalUnitTestBase {
   }
 
   /**
-   * Tests setPaymentContext() and getPaymentContext().
+   * Tests bundle().
+   */
+  function bundle() {
+    $this->assertIdentical($this->payment->bundle(), $this->bundle);
+  }
+
+  /**
+   * Tests getPaymentContext().
    */
   function testGetPaymentContext() {
-    $context = \Drupal::service('plugin.manager.payment.context')->createInstance('payment_unavailable');
-    $this->assertTrue($this->payment->setPaymentContext($context) instanceof PaymentInterface);
-    $this->assertIdentical($this->payment->getPaymentContext(), $context);
+    if ($this->assertTrue($this->payment->getPaymentContext() instanceof PaymentContextInterface)) {
+      $this->assertIdentical($this->payment->getPaymentContext()->getPluginId(), $this->bundle);
+    }
   }
 
   /**
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/payment/context/UnavailableUnitTest.php b/payment/lib/Drupal/payment/Tests/Plugin/payment/context/UnavailableUnitTest.php
index ba64cf3..6b78d8b 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/payment/context/UnavailableUnitTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/payment/context/UnavailableUnitTest.php
@@ -65,7 +65,9 @@ class UnavailableUnitTest extends DrupalUnitTestBase {
    * Tests setPayment() and getPayment().
    */
   function testGetPayment() {
-    $payment = entity_create('payment', array());
+    $payment = entity_create('payment', array(
+      'context' => $this->context,
+    ));
     $this->assertTrue($this->context->setPayment($payment) instanceof PaymentContextInterface );
     $this->assertTrue($this->context->getPayment() === $payment);
   }
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/payment/line_item/BasicWebTest.php b/payment/lib/Drupal/payment/Tests/Plugin/payment/line_item/BasicWebTest.php
index 6551b29..b84b23e 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/payment/line_item/BasicWebTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/payment/line_item/BasicWebTest.php
@@ -95,7 +95,9 @@ class BasicWebTest extends WebTestBase {
   public function testPaymentEntityIntegration() {
     $description = $this->randomString();
     $name = $this->randomName();
-    $payment = entity_create('payment', array());
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ));
     $this->lineItem
       ->setName($name)
       ->setDescription($description);
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicUnitTest.php b/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicUnitTest.php
index f9268a3..34fbb46 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicUnitTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicUnitTest.php
@@ -117,7 +117,9 @@ class BasicUnitTest extends DrupalUnitTestBase {
    * Tests paymentOperationAccess().
    */
   function testPaymentOperationAccess() {
-    $payment = entity_create('payment', array());
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ));
     $this->assertTrue($this->method->paymentOperationAccess($payment, 'execute', 'default'));
     $this->assertFalse($this->method->paymentOperationAccess($payment, $this->randomName(), 'default'));
   }
@@ -127,7 +129,9 @@ class BasicUnitTest extends DrupalUnitTestBase {
    */
   function testExecutePaymentOperation() {
     $plugin_id = 'payment_unknown';
-    $payment = entity_create('payment', array());
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ));
     $this->method->setStatus($plugin_id);
     $this->method->executePaymentOperation($payment, 'execute', 'default');
     $this->assertTrue($payment->getStatus()->getPluginId() == $plugin_id);
diff --git a/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicWebTest.php b/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicWebTest.php
index c9b0658..c3986b0 100644
--- a/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicWebTest.php
+++ b/payment/lib/Drupal/payment/Tests/Plugin/payment/method/BasicWebTest.php
@@ -43,10 +43,10 @@ class BasicWebTest extends WebTestBase {
   function testPaymentFormElements() {
     $this->method->setMessageText('Hello [site:name]!');
     $form = array();
-    $form_state = array(
-      'payment' => entity_create('payment', array()),
-    );
-    $payment = entity_create('payment', array());
+    $form_state = array();
+    $payment = entity_create('payment', array(
+      'bundle' => 'payment_unavailable',
+    ));
     $elements = $this->method->paymentFormElements($form, $form_state, $payment);
     if ($this->assertTrue(is_array($elements))) {
       $this->assertIdentical($elements['message']['#markup'], "<p>Hello Drupal!</p>\n");
diff --git a/payment/payment.install b/payment/payment.install
index 81c293c..f8bd6e2 100644
--- a/payment/payment.install
+++ b/payment/payment.install
@@ -102,7 +102,7 @@ function payment_schema() {
   );
   $schema['payment'] = array(
     'fields' => array(
-      'context_plugin_id' => array(
+      'bundle' => array(
         'type' => 'varchar',
         'length' => 255,
       ),
@@ -424,8 +424,8 @@ function payment_update_8200(array &$sandbox) {
       ->execute();
   }
 
-  // Rename context to context_plugin_id.
-  db_change_field('payment', 'context', 'context_plugin_id', array(
+  // Rename context to bundle.
+  db_change_field('payment', 'context', 'bundle', array(
     'type' => 'varchar',
     'length' => 255,
   ));
diff --git a/payment_test/lib/Drupal/payment_test/PaymentMethodElement.php b/payment_test/lib/Drupal/payment_test/PaymentMethodElement.php
index 3e215dc..0310d88 100644
--- a/payment_test/lib/Drupal/payment_test/PaymentMethodElement.php
+++ b/payment_test/lib/Drupal/payment_test/PaymentMethodElement.php
@@ -57,9 +57,9 @@ class PaymentMethodElement implements ControllerInterface, FormInterface {
    * {@inheritdoc}
    */
   public function buildForm(array $form, array &$form_state) {
-    $payment = $this->entityManager->getStorageController('payment')->create(array())
-      ->setPaymentContext($this->contextManager->createInstance('payment_unavailable'))
-      ->setLineItems(Generate::createPaymentLineItems());
+    $payment = $this->entityManager->getStorageController('payment')->create(array(
+      'bundle' => 'payment_unavailable',
+    ))->setLineItems(Generate::createPaymentLineItems());
     $form['payment_method'] = array(
       '#default_value' => $payment,
       '#required' => TRUE,