From b790bc9fc6ba8d060c473fb3020220157ae72ac4 Mon Sep 17 00:00:00 2001
From: Michal Proszek <michal.proszek@gmail.com>
Date: Wed, 6 Jun 2018 20:55:03 +0200
Subject: [PATCH] [BE] optimize refresh_customers command

---
 CHANGELOG.txt                                 |  4 ++
 aa_stripe/__init__.py                         |  2 +-
 .../management/commands/refresh_customers.py  | 37 ++++++++++++--
 tests/test_customers.py                       | 51 ++++++++++++++++---
 4 files changed, 82 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.txt b/CHANGELOG.txt
index aaf25c2..698bcdb 100644
--- a/CHANGELOG.txt
+++ b/CHANGELOG.txt
@@ -1,5 +1,9 @@
 # Change Log
 All notable changes to this project will be documented in this file.
+## [0.6.4]
+### Added
+- refresh_customers command optimization
+
 ## [0.6.3]
 ### Fixed
 - saving StripeCustomer.default_source
diff --git a/aa_stripe/__init__.py b/aa_stripe/__init__.py
index 3c1f03f..1586682 100644
--- a/aa_stripe/__init__.py
+++ b/aa_stripe/__init__.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 __title__ = "Arabella Stripe"
-__version__ = "0.6.3"
+__version__ = "0.6.4"
 __author__ = "Jacek Ostanski"
 __license__ = "MIT"
 __copyright__ = "Copyright 2017 Arabella"
diff --git a/aa_stripe/management/commands/refresh_customers.py b/aa_stripe/management/commands/refresh_customers.py
index dabdfd9..e43440e 100644
--- a/aa_stripe/management/commands/refresh_customers.py
+++ b/aa_stripe/management/commands/refresh_customers.py
@@ -1,6 +1,9 @@
 # -*- coding: utf-8 -*-
+import sys
+
 import stripe
 from django.core.management.base import BaseCommand
+from django.utils import timezone
 
 from aa_stripe.models import StripeCustomer
 from aa_stripe.settings import stripe_settings
@@ -11,8 +14,36 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         stripe.api_key = stripe_settings.API_KEY
-        for customer in StripeCustomer.objects.filter(is_active=True, is_created_at_stripe=True):
+        last_customer = None
+        retry_count = 0
+        updated_count = 0
+        start_time = timezone.now()
+        print("Began refreshing customers")
+        while True:
             try:
-                customer.refresh_from_stripe()
+                response = stripe.Customer.list(limit=100, starting_after=last_customer)
             except stripe.StripeError:
-                pass
+                if retry_count > 5:
+                    raise
+                retry_count += 1
+                continue
+            else:
+                retry_count = 0
+
+            for stripe_customer in response["data"]:
+                updated_count += StripeCustomer.objects.filter(stripe_customer_id=stripe_customer["id"]).update(
+                    sources=stripe_customer["sources"], default_source=stripe_customer["default_source"]
+                )
+
+            if not response["has_more"]:
+                break
+
+            sys.stdout.write(".")  # indicate that the command did not hung up
+            sys.stdout.flush()
+            last_customer = response["data"][-1]
+
+        if updated_count:
+            print("\nCustomers updated: {} (took {:2f}s)".format(
+                updated_count, (timezone.now() - start_time).total_seconds()))
+        else:
+            print("No customers were updated.")
diff --git a/tests/test_customers.py b/tests/test_customers.py
index 7cdac56..6b822a5 100644
--- a/tests/test_customers.py
+++ b/tests/test_customers.py
@@ -142,16 +142,51 @@ def test_change_description(self):
 class TestRefreshCustomersCommand(BaseTestCase):
     def setUp(self):
         self._create_customer(is_active=False, is_created_at_stripe=False)
-        self._create_customer()
-
-    @mock.patch("aa_stripe.models.StripeCustomer.refresh_from_stripe")
-    def test_command(self, mocked_update):
-        call_command("refresh_customers")
-        mocked_update.assert_called_once()
+        self.active_customer = self._create_customer()
 
-        # the command should not exit in case of an error during update (for example customer does not exist)
-        mocked_update.side_effect = stripe.APIError
+    @requests_mock.Mocker()
+    def test_command(self, m):
+        customer1_data = {
+            "id": "cus_xyz",
+            "object": "customer",
+            "sources": [
+                {"id": "card_1"}
+            ],
+            "default_source": "card_1"
+        }
+        customer2_data = customer1_data.copy()
+        customer2_data["id"] = "cus_2"
+        customer2_data["sources"] = [{"id": "card_2"}]
+        customer2_data["default_source"] = "card_2"
+        customer2_data = customer1_data.copy()
+        customer2_data["id"] = "cus_2"
+        customer2_data["sources"] = []
+        customer2_data["default_source"] = None
+        stripe_response_part1 = {
+            "object": "list",
+            "url": "/v1/customers",
+            "has_more": False,
+            "data": [
+                customer1_data
+            ]
+        }
+        stripe_response_part2 = stripe_response_part1.copy()
+        stripe_response_part2["has_more"] = False
+        stripe_response_part2["data"] = [customer2_data]
+        m.register_uri("GET", "https://api.stripe.com/v1/customers", text=json.dumps(stripe_response_part1))
+        m.register_uri("GET", "https://api.stripe.com/v1/customers?starting_after=cus_xyz", [
+            {"text": "", "status_code": 500},  # make sure the command will try again
+            {"text": json.dumps(stripe_response_part2), "status_code": "200"}
+        ])
         call_command("refresh_customers")
+        self.active_customer.refresh_from_db()
+        self.assertEqual(self.active_customer.get_default_source_data(), {"id": "card_1"})
+
+        # the command should fail if call to api fails more than 5 times
+        with mock.patch("stripe.Customer.list") as mocked_list:
+            mocked_list.side_effect = stripe.APIError()
+            with self.assertRaises(stripe.APIError):
+                call_command("refresh_customers")
 
     @requests_mock.Mocker()
     def test_customer_refresh_from_stripe(self, m):