Add fixes for PR dj-stripe#320 review

Adds a new _attach_objects_post_save_hook, for attaching Stripe objects
to an instance only after it is saved.

djstripe/models.py

@@ -346,6 +346,16 @@ def add_card(self, source, set_default=True):
 
         return new_card
 
+    def upcoming_invoice(self, **kwargs):
+        """ Gets the upcoming preview invoice (singular) for this customer.
+
+        See `StripeInvoice.upcoming() <#djstripe.stripe_objects.StripeInvoice.upcoming>`__
+        The ``customer`` argument to the ``upcoming()`` call is automatically set by this method.
+        """
+
+        kwargs['customer'] = self
+        return Invoice.upcoming(**kwargs)
+
     def _attach_objects_hook(self, cls, data):
         # TODO: other sources
         if data["default_source"] and data["default_source"]["object"] == "card":
@@ -542,35 +552,41 @@ def _attach_objects_hook(self, cls, data):
         if subscription:
             self.subscription = subscription
 
-        self._attach_invoice_items(cls, data)
+    def _attach_objects_post_save_hook(self, cls, data):
+        # InvoiceItems need a saved invoice because they're associated via a
+        # RelatedManager, so this must be done as part of the post save hook.
+        cls._stripe_object_to_invoice_items(target_cls=InvoiceItem, data=data, invoice=self)
 
-    def _attach_invoice_items(self, cls, data):
-        if not self.pk:
-            # InvoiceItems need a saved invoice because they're associated via
-            # a RelatedManager.  It might be better to make this pattern more
-            # generic with some sort-of _attach_objects_post_save_hook().
-            self.save()
+    @classmethod
+    def upcoming(cls, **kwargs):
+        upcoming_stripe_invoice = StripeInvoice.upcoming(**kwargs)
 
-        cls._stripe_object_to_invoice_items(InvoiceItem, data, self)
+        if upcoming_stripe_invoice:
+            return UpcomingInvoice._create_from_stripe_object(upcoming_stripe_invoice, save=False)
 
     @property
     def plan(self):
         """ Gets the associated plan for this invoice.
 
-        In order to provide a consistent view of invoices, it is necessary to
-        examine the invoice items for the associated subscription plan rather
-        than the top-level link.  The reason for this is that the subscription
-        plan when requested by the customer, but the invoice item plan will
-        remain the same.  This makes it difficult to work with an invoice
-        history if the plan is expected to remain consistent (e.g. for creating
-        downloadable invoices).
+        In order to provide a consistent view of invoices, the plan object
+        should be taken from the first invoice item that has one, rather than
+        using the plan associated with the subscription.
+
+        Subscriptions (and their associated plan) are updated by the customer
+        and represent what is current, but invoice items are immutable within
+        the invoice and stay static/unchanged.
+
+        In other words, a plan retrieved from an invoice item will represent
+        the plan as it was at the time an invoice was issued.  The plan
+        retrieved from the subscription will be the currently active plan.
 
         :returns: The associated plan for the invoice.
         :rtype: ``djstripe.models.Plan``
         """
-        for item in self.invoiceitems.all():
-            if item.plan:
-                return item.plan
+
+        for invoiceitem in self.invoiceitems.all():
+            if invoiceitem.plan:
+                return invoiceitem.plan
 
         if self.subscription:
             return self.subscription.plan
@@ -582,24 +598,26 @@ class UpcomingInvoice(Invoice):
 
     def __init__(self, *args, **kwargs):
         super(UpcomingInvoice, self).__init__(*args, **kwargs)
-        self._items = []
+        self._invoiceitems = []
 
-    def _attach_invoice_items(self, cls, data):
-        self._items = cls._stripe_object_to_invoice_items(InvoiceItem, data, self)
+    def _attach_objects_hook(self, cls, data):
+        super(UpcomingInvoice, self)._attach_objects_hook(cls, data)
+        self._invoiceitems = cls._stripe_object_to_invoice_items(target_cls=InvoiceItem, data=data, invoice=self)
 
     @property
     def invoiceitems(self):
         """ Gets the invoice items associated with this upcoming invoice.
 
         This differs from normal (non-upcoming) invoices, in that upcoming
-        invoices are in-memory and do not persist to the database.  Therefore,
+        invoices are in-memory and do not persist to the database. Therefore,
         all of the data comes from the Stripe API itself.
 
         Instead of returning a normal queryset for the invoiceitems, this will
         return a mock of a queryset, but with the data fetched from Stripe - It
         will act like a normal queryset, but mutation will silently fail.
         """
-        return QuerySetMock(InvoiceItem, *self._items)
+
+        return QuerySetMock(InvoiceItem, *self._invoiceitems)
 
     @property
     def stripe_id(self):

djstripe/stripe_objects.py

@@ -24,7 +24,7 @@
 
 from django.conf import settings
 from django.db import models
-from django.utils import six, timezone
+from django.utils import dateformat, six, timezone
 from django.utils.encoding import python_2_unicode_compatible, smart_text
 from model_utils.models import TimeStampedModel
 from polymorphic.models import PolymorphicModel
@@ -176,6 +176,22 @@ def _attach_objects_hook(self, cls, data):
         """
         Gets called by this object's create and sync methods just before save.
         Use this to populate fields before the model is saved.
+
+        :param cls: The target class for the instantiated object.
+        :param data: The data dictionary received from the Stripe API.
+        :type data: dict
+        """
+
+        pass
+
+    def _attach_objects_post_save_hook(self, cls, data):
+        """
+        Gets called by this object's create and sync methods just after save.
+        Use this to populate fields after the model is saved.
+
+        :param cls: The target class for the instantiated object.
+        :param data: The data dictionary received from the Stripe API.
+        :type data: dict
         """
 
         pass
@@ -185,21 +201,26 @@ def _create_from_stripe_object(cls, data, save=True):
         """
         Instantiates a model instance using the provided data object received
         from Stripe, and saves it to the database if specified.
+
         :param data: The data dictionary received from the Stripe API.
         :type data: dict
         :param save: If True, the object is saved after instantiation.
         :type save: bool
         :returns: The instantiated object.
         """
+
         instance = cls(**cls._stripe_object_to_record(data))
         instance._attach_objects_hook(cls, data)
+
         if save:
             instance.save()
+
+        instance._attach_objects_post_save_hook(cls, data)
+
         return instance
 
     @classmethod
-    def _get_or_create_from_stripe_object(cls, data, field_name="id",
-                                          refetch=True, save=True):
+    def _get_or_create_from_stripe_object(cls, data, field_name="id", refetch=True, save=True):
         field = data.get(field_name)
 
         if isinstance(field, six.string_types):
@@ -291,18 +312,21 @@ def _stripe_object_to_invoice_items(cls, target_cls, data, invoice):
         If the invoice item doesn't exist already then it is created.
 
         If the invoice is an upcoming invoice that doesn't persist to the
-        database (i.e. ephermeral) then the invoice items are also not saved.
+        database (i.e. ephemeral) then the invoice items are also not saved.
 
         :param target_cls: The target class to instantiate per invoice item.
         :type target_cls: ``StripeInvoiceItem``
         :param data: The data dictionary received from the Stripe API.
         :type data: dict
+        :param invoice: The invoice object that should hold the invoice items.
+        :type invoice: ``djstripe.models.Invoice``
         """
+
         lines = data.get("lines")
         if not lines:
             return []
 
-        items = []
+        invoiceitems = []
         for line in lines.get("data", []):
             if invoice.stripe_id:
                 save = True
@@ -311,25 +335,25 @@ def _stripe_object_to_invoice_items(cls, target_cls, data, invoice):
                 if line.get("type") == "subscription":
                     # Lines for subscriptions need to be keyed based on invoice and
                     # subscription, because their id is *just* the subscription
-                    # when received from Stripe.  This means that future updates to
+                    # when received from Stripe. This means that future updates to
                     # a subscription will change previously saved invoices - Doing
                     # the composite key avoids this.
                     if not line["id"].startswith(invoice.stripe_id):
-                        line["id"] = "{invoice_id}-{sub_id}".format(
+                        line["id"] = "{invoice_id}-{subscription_id}".format(
                             invoice_id=invoice.stripe_id,
-                            sub_id=line["id"])
+                            subscription_id=line["id"])
             else:
                 # Don't save invoice items for ephemeral invoices
                 save = False
 
             line.setdefault("customer", invoice.customer.stripe_id)
-            line.setdefault("date", int(invoice.date.strftime("%s")))
+            line.setdefault("date", int(dateformat.format(invoice.date, 'U')))
 
-            item = target_cls._get_or_create_from_stripe_object(
-                line, refetch=False, save=save)[0]
-            items.append(item)
+            item, _ = target_cls._get_or_create_from_stripe_object(
+                line, refetch=False, save=save)
+            invoiceitems.append(item)
 
-        return items
+        return invoiceitems
 
     @classmethod
     def _stripe_object_to_subscription(cls, target_cls, data):
@@ -364,6 +388,7 @@ def sync_from_stripe_data(cls, data):
             instance._sync(cls._stripe_object_to_record(data))
             instance._attach_objects_hook(cls, data)
             instance.save()
+            instance._attach_objects_post_save_hook(cls, data)
 
         return instance
 
@@ -1168,7 +1193,9 @@ def _stripe_object_to_charge(cls, target_cls, data):
             return target_cls._get_or_create_from_stripe_object(data, "charge")[0]
 
     @classmethod
-    def upcoming(cls, api_key=settings.STRIPE_SECRET_KEY, **kwargs):
+    def upcoming(cls, api_key=settings.STRIPE_SECRET_KEY, customer=None, coupon=None, subscription=None,
+                 subscription_plan=None, subscription_prorate=None, subscription_proration_date=None,
+                 subscription_quantity=None, subscription_trial_end=None, **kwargs):
         """
         Gets the upcoming preview invoice (singular) for a customer.
 
@@ -1182,45 +1209,52 @@ def upcoming(cls, api_key=settings.STRIPE_SECRET_KEY, **kwargs):
 
         :param customer: The identifier of the customer whose upcoming invoice
         you'd like to retrieve.
+        :type customer: ``djstripe.models.Customer``
         :param coupon: The code of the coupon to apply.
+        :type customer: ``str``
         :param subscription: The identifier of the subscription to retrieve an
         invoice for.
+        :type customer: ``djstripe.models.Subscription`` or ``str`` (id)
         :param subscription_plan: If set, the invoice returned will preview
         updating the subscription given to this plan, or creating a new
         subscription to this plan if no subscription is given.
+        :type subscription_plan: ``djstripe.models.Subscription`` or ``str`` (id)
         :param subscription_prorate: If previewing an update to a subscription,
         this decides whether the preview will show the result of applying
         prorations or not.
-        :param subscription_proration_date: f previewing an update to a
+        :type subscription_prorate: ``bool``
+        :param subscription_proration_date: If previewing an update to a
         subscription, and doing proration, subscription_proration_date forces
         the proration to be calculated as though the update was done at the
         specified time.
+        :type subscription_proration_date: ``datetime``
         :param subscription_quantity: If provided, the invoice returned will
         preview updating or creating a subscription with that quantity.
+        :type subscription_proration_quantity: ``int``
         :param subscription_trial_end: If provided, the invoice returned will
         preview updating or creating a subscription with that trial end.
+        :type subscription_trial_end: ``datetime``
         :returns: The upcoming preview invoice.
         :rtype: ``djstripe.models.UpcomingInvoice``
         """
-        from . models import Invoice, UpcomingInvoice
-        if not issubclass(cls, UpcomingInvoice):
-            if cls is Invoice:
-                cls = UpcomingInvoice
-            else:
-                assert False, (
-                    "Class argument needs to be derived from UpcomingInvoice")
-
         try:
-            data = cls._api().upcoming(api_key=api_key, **kwargs)
-        except InvalidRequestError as ex:
-            if str(ex) != "Nothing to invoice for customer":
+            upcoming_stripe_invoice = cls._api().upcoming(
+                api_key=api_key, customer=customer,
+                coupon=coupon, subscription=subscription,
+                subscription_plan=subscription_plan,
+                subscription_prorate=subscription_prorate,
+                subscription_proration_date=subscription_proration_date,
+                subscription_quantity=subscription_quantity,
+                subscription_trial_end=subscription_trial_end, **kwargs)
+        except InvalidRequestError as exc:
+            if str(exc) != "Nothing to invoice for customer":
                 six.reraise(*sys.exc_info())
             return
 
         # Workaround for "id" being missing (upcoming invoices don't persist).
-        data["id"] = "upcoming"
+        upcoming_stripe_invoice["id"] = "upcoming"
 
-        return cls._create_from_stripe_object(data, save=False)
+        return upcoming_stripe_invoice
 
     def retry(self):
         """ Retry payment on this invoice if it isn't paid, closed, or forgiven."""

docs/models.rst

@@ -51,6 +51,7 @@ Customer
     .. automethod:: djstripe.models.Customer.retry_unpaid_invoices
     .. automethod:: djstripe.models.Customer.has_valid_source
     .. automethod:: djstripe.models.Customer.add_card
+    .. automethod:: djstripe.models.Customer.upcoming_invoice
 
     .. automethod:: djstripe.models.Customer.str_parts
     .. automethod:: djstripe.stripe_objects.StripeObject.sync_from_stripe_data
@@ -123,8 +124,10 @@ Invoice
     .. autoattribute:: djstripe.models.Invoice.STATUS_CLOSED
     .. autoattribute:: djstripe.models.Invoice.STATUS_OPEN
     .. autoattribute:: djstripe.models.Invoice.status
+    .. autoattribute:: djstripe.models.Invoice.plan
 
     .. automethod:: djstripe.models.Invoice.retry
+    .. automethod:: djstripe.models.Invoice.upcoming
 
     .. automethod:: djstripe.models.Invoice.str_parts
     .. automethod:: djstripe.stripe_objects.StripeObject.sync_from_stripe_data

tests/test_customer.py

@@ -5,6 +5,7 @@
 .. moduleauthor:: Daniel Greenfeld (@pydanny)
 .. moduleauthor:: Alex Kavanaugh (@kavdev)
 .. moduleauthor:: Michael Thornhill (@mthornhill)
+.. moduleauthor:: Lee Skillen (@lskillen)
 
 """
 
@@ -16,14 +17,14 @@
 from django.contrib.auth import get_user_model
 from django.test import TestCase
 from django.utils import timezone
-from mock import patch
+from mock import patch, ANY
 from stripe.error import InvalidRequestError
 
 from djstripe.exceptions import MultipleSubscriptionException
 from djstripe.models import Account, Customer, Charge, Card, Subscription, Invoice, Plan
 from tests import (FAKE_CARD, FAKE_CHARGE, FAKE_CUSTOMER, FAKE_ACCOUNT, FAKE_INVOICE,
                    FAKE_INVOICE_III, FAKE_INVOICEITEM, FAKE_PLAN, FAKE_SUBSCRIPTION, FAKE_SUBSCRIPTION_II,
-                   StripeList, FAKE_CARD_V, FAKE_CUSTOMER_II)
+                   StripeList, FAKE_CARD_V, FAKE_CUSTOMER_II, FAKE_UPCOMING_INVOICE)
 
 
 class TestCustomer(TestCase):
@@ -576,3 +577,27 @@ def test_add_invoice_item_djstripe_objects(self, invoiceitem_create_mock, invoic
     def test_add_invoice_item_bad_decimal(self):
         with self.assertRaisesMessage(ValueError, "You must supply a decimal value representing dollars."):
             self.customer.add_invoice_item(amount=5000, currency="usd")
+
+    @patch("stripe.Plan.retrieve", return_value=deepcopy(FAKE_PLAN))
+    @patch("stripe.Subscription.retrieve", return_value=deepcopy(FAKE_SUBSCRIPTION))
+    @patch("stripe.Invoice.upcoming", return_value=deepcopy(FAKE_UPCOMING_INVOICE))
+    def test_upcoming_invoice(self, invoice_upcoming_mock, subscription_retrieve_mock, plan_retrieve_mock):
+        invoice = self.customer.upcoming_invoice()
+        self.assertIsNotNone(invoice)
+        self.assertIsNone(invoice.stripe_id)
+        self.assertIsNone(invoice.save())
+
+        subscription_retrieve_mock.assert_called_once_with(api_key=ANY, expand=ANY, id=FAKE_SUBSCRIPTION["id"])
+        plan_retrieve_mock.assert_not_called()
+
+        items = invoice.invoiceitems.all()
+        self.assertEquals(1, len(items))
+        self.assertEquals(FAKE_SUBSCRIPTION["id"], items[0].stripe_id)
+
+        self.assertIsNotNone(invoice.plan)
+        self.assertEquals(FAKE_PLAN["id"], invoice.plan.stripe_id)
+
+        invoice._invoiceitems = []
+        items = invoice.invoiceitems.all()
+        self.assertEquals(0, len(items))
+        self.assertIsNotNone(invoice.plan)