Adding tests to hc.payments

hc/front/views.py

@@ -278,6 +278,7 @@ def add_channel(request):
     assert request.method == "POST"
     return do_add_channel(request, request.POST)
 
+
 @login_required
 @uuid_or_400
 def channel_checks(request, code):

hc/payments/models.py

@@ -1,8 +1,19 @@
-import braintree
 from django.contrib.auth.models import User
 from django.db import models
 
 
+class SubscriptionManager(models.Manager):
+
+    def for_user(self, user):
+        try:
+            sub = self.get(user_id=user.id)
+        except Subscription.DoesNotExist:
+            sub = Subscription(user=user)
+            sub.save()
+
+        return sub
+
+
 class Subscription(models.Model):
     user = models.OneToOneField(User, blank=True, null=True)
     customer_id = models.CharField(max_length=36, blank=True)
@@ -10,24 +21,7 @@ class Subscription(models.Model):
     subscription_id = models.CharField(max_length=10, blank=True)
     plan_id = models.CharField(max_length=10, blank=True)
 
-    def _get_braintree_sub(self):
+    objects = SubscriptionManager()
-        if not hasattr(self, "_sub"):
-            self._sub = braintree.Subscription.find(self.subscription_id)
-
-        return self._sub
-
-    def _get_braintree_payment_method(self):
-        if not hasattr(self, "_pm"):
-            self._pm = braintree.PaymentMethod.find(self.payment_method_token)
-
-        return self._pm
-
-    def is_active(self):
-        if not self.subscription_id:
-            return False
-
-        o = self._get_braintree_sub()
-        return o.status == "Active"
 
     def price(self):
         if self.plan_id == "P5":
@@ -36,27 +30,3 @@ class Subscription(models.Model):
             return 20
 
         return 0
-
-    def next_billing_date(self):
-        o = self._get_braintree_sub()
-        return o.next_billing_date
-
-    def pm_is_credit_card(self):
-        return isinstance(self._get_braintree_payment_method(),
-                          braintree.credit_card.CreditCard)
-
-    def pm_is_paypal(self):
-        return isinstance(self._get_braintree_payment_method(),
-                          braintree.paypal_account.PayPalAccount)
-
-    def card_type(self):
-        o = self._get_braintree_payment_method()
-        return o.card_type
-
-    def last_4(self):
-        o = self._get_braintree_payment_method()
-        return o.last_4
-
-    def paypal_email(self):
-        o = self._get_braintree_payment_method()
-        return o.email

hc/payments/tests.py

@@ -1,3 +0,0 @@
-from django.test import TestCase
-
-# Create your tests here.

hc/payments/tests/test_create_plan.py

@@ -0,0 +1,67 @@
+from django.contrib.auth.models import User
+from django.test import TestCase
+from hc.accounts.models import Profile
+from hc.payments.models import Subscription
+from mock import patch
+
+
+class CreatePlanTestCase(TestCase):
+
+    def setUp(self):
+        self.alice = User(username="alice")
+        self.alice.set_password("password")
+        self.alice.save()
+
+    def _setup_mock(self, mock):
+        """ Set up Braintree calls that the controller will use. """
+
+        mock.Customer.create.return_value.is_success = True
+        mock.Customer.create.return_value.customer.id = "test-customer-id"
+
+        mock.PaymentMethod.create.return_value.is_success = True
+        mock.PaymentMethod.create.return_value.payment_method.token = "t-token"
+
+        mock.Subscription.create.return_value.is_success = True
+        mock.Subscription.create.return_value.subscription.id = "t-sub-id"
+
+    def run_create_plan(self, plan_id="P5"):
+        form = {"plan_id": plan_id, "payment_method_nonce": "test-nonce"}
+        self.client.login(username="alice", password="password")
+        return self.client.post("/pricing/create_plan/", form)
+
+    @patch("hc.payments.views.braintree")
+    def test_it_works(self, mock):
+        self._setup_mock(mock)
+
+        r = self.run_create_plan()
+        self.assertEqual(r.status_code, 302)
+
+        # Subscription should be filled out:
+        sub = Subscription.objects.get(user=self.alice)
+        self.assertEqual(sub.customer_id, "test-customer-id")
+        self.assertEqual(sub.payment_method_token, "t-token")
+        self.assertEqual(sub.subscription_id, "t-sub-id")
+        self.assertEqual(sub.plan_id, "P5")
+
+        # User's profile should have a higher ping log limit:
+        profile = Profile.objects.get(user=self.alice)
+        self.assertEqual(profile.ping_log_limit, 1000)
+
+        # braintree.Subscription.cancel should have not been called
+        assert not mock.Subscription.cancel.called
+
+    def test_bad_plan_id(self):
+        r = self.run_create_plan(plan_id="this-is-wrong")
+        self.assertEqual(r.status_code, 400)
+
+    @patch("hc.payments.views.braintree")
+    def test_it_cancels_previous_subscription(self, mock):
+        self._setup_mock(mock)
+
+        sub = Subscription(user=self.alice)
+        sub.subscription_id = "prev-sub"
+        sub.save()
+
+        r = self.run_create_plan()
+        self.assertEqual(r.status_code, 302)
+        assert mock.Subscription.cancel.called

hc/payments/tests/test_get_client_token.py

@@ -0,0 +1,23 @@
+from django.contrib.auth.models import User
+from django.test import TestCase
+from hc.payments.models import Subscription
+from mock import patch
+
+
+class GetClientTokenTestCase(TestCase):
+
+    def setUp(self):
+        self.alice = User(username="alice")
+        self.alice.set_password("password")
+        self.alice.save()
+
+    @patch("hc.payments.views.braintree")
+    def test_it_works(self, mock_braintree):
+        mock_braintree.ClientToken.generate.return_value = "test-token"
+        self.client.login(username="alice", password="password")
+
+        r = self.client.get("/pricing/get_client_token/")
+        self.assertContains(r, "test-token", status_code=200)
+
+        # A subscription object should have been created
+        assert Subscription.objects.count() == 1

hc/payments/tests/test_pricing.py

@@ -0,0 +1,27 @@
+from django.contrib.auth.models import User
+from django.test import TestCase
+from hc.payments.models import Subscription
+
+
+class PricingTestCase(TestCase):
+
+    def setUp(self):
+        self.alice = User(username="alice")
+        self.alice.set_password("password")
+        self.alice.save()
+
+    def test_anonymous(self):
+        r = self.client.get("/pricing/")
+        self.assertContains(r, "Unlimited Checks", status_code=200)
+
+        # A subscription object should have NOT been created
+        assert Subscription.objects.count() == 0
+
+    def test_authenticated(self):
+        self.client.login(username="alice", password="password")
+
+        r = self.client.get("/pricing/")
+        self.assertContains(r, "Unlimited Checks", status_code=200)
+
+        # A subscription object should have been created
+        assert Subscription.objects.count() == 1

hc/payments/views.py

@@ -1,17 +1,18 @@
 import braintree
-from django.contrib.auth.decorators import login_required
 from django.contrib import messages
-from django.http import HttpResponseForbidden, JsonResponse
+from django.contrib.auth.decorators import login_required
+from django.http import (HttpResponseBadRequest, HttpResponseForbidden,
+                         JsonResponse)
 from django.shortcuts import redirect, render
 from django.views.decorators.http import require_POST
-
 from hc.accounts.models import Profile
+
 from .models import Subscription
 
 
 @login_required
 def get_client_token(request):
-    sub = Subscription.objects.get(user=request.user)
+    sub = Subscription.objects.for_user(request.user)
     client_token = braintree.ClientToken.generate({
         "customer_id": sub.customer_id
     })
@@ -22,21 +23,12 @@ def get_client_token(request):
 def pricing(request):
     sub = None
     if request.user.is_authenticated():
-        try:
+        sub = Subscription.objects.for_user(request.user)
-            sub = Subscription.objects.get(user=request.user)
-        except Subscription.DoesNotExist:
-            sub = Subscription(user=request.user)
-            sub.save()
-
-    first_charge = False
-    if "first_charge" in request.session:
-        first_charge = True
-        del request.session["first_charge"]
 
     ctx = {
         "page": "pricing",
         "sub": sub,
-        "first_charge": first_charge
+        "first_charge": request.session.pop("first_charge", False)
     }
 
     return render(request, "payments/pricing.html", ctx)
@@ -55,9 +47,10 @@ def log_and_bail(request, result):
 @require_POST
 def create_plan(request):
     plan_id = request.POST["plan_id"]
-    assert plan_id in ("P5", "P20")
+    if plan_id not in ("P5", "P20"):
+        return HttpResponseBadRequest()
 
-    sub = Subscription.objects.get(user=request.user)
+    sub = Subscription.objects.for_user(request.user)
 
     # Cancel the previous plan
     if sub.subscription_id:
@@ -133,9 +126,10 @@ def cancel_plan(request):
 def billing(request):
     sub = Subscription.objects.get(user=request.user)
 
-    transactions = braintree.Transaction.search(braintree.TransactionSearch.customer_id == sub.customer_id)
+    transactions = braintree.Transaction.search(
-    ctx = {"transactions": transactions}
+        braintree.TransactionSearch.customer_id == sub.customer_id)
 
+    ctx = {"transactions": transactions}
     return render(request, "payments/billing.html", ctx)