write tests for low balance email check, add indexes, reduce cuttoff …

…to 200 credits

avoid hardcoding urls and dates

app_users/models.py

@@ -265,6 +265,10 @@ class AppUserTransaction(models.Model):
 
     class Meta:
         verbose_name = "Transaction"
+        indexes = [
+            models.Index(fields=["user", "amount", "-created_at"]),
+            models.Index(fields=["-created_at"]),
+        ]
 
     def __str__(self):
         return f"{self.invoice_id} ({self.amount})"

celeryapp/tasks.py

@@ -1,13 +1,15 @@
-from fastapi import HTTPException
+import datetime
 import html
 import traceback
 import typing
-import datetime
 from time import time
 from types import SimpleNamespace
 
 import requests
 import sentry_sdk
+from django.db.models import Sum
+from django.utils import timezone
+from fastapi import HTTPException
 
 import gooey_ui as st
 from app_users.models import AppUser, AppUserTransaction
@@ -18,13 +20,11 @@
 from daras_ai_v2.base import StateKeys, BasePage
 from daras_ai_v2.exceptions import UserError
 from daras_ai_v2.send_email import send_email_via_postmark
+from daras_ai_v2.send_email import send_low_balance_email
 from daras_ai_v2.settings import templates
 from gooey_ui.pubsub import realtime_push
 from gooey_ui.state import set_query_params
 from gooeysite.bg_db_conn import db_middleware, next_db_safe
-from daras_ai_v2.send_email import send_low_balance_email
-from django.db.models import Sum
-from django.utils import timezone
 
 
 @app.task
@@ -128,9 +128,9 @@ def save(done=False):
                 save()
     finally:
         save(done=True)
-        low_balance_email(sr)
         if not is_api_call:
             send_email_on_completion(page, sr)
+        run_low_balance_email_check(uid)
 
 
 def err_msg_for_exc(e: Exception):
@@ -158,42 +158,41 @@ def err_msg_for_exc(e: Exception):
         return f"{type(e).__name__}: {e}"
 
 
-def low_balance_email(sr: SavedRun):
-    user = AppUser.objects.get(uid=sr.uid)
-    last_positive_transaction = (
+def run_low_balance_email_check(uid: str):
+    # don't send email if feature is disabled
+    if not settings.LOW_BALANCE_EMAIL_ENABLED:
+        return
+    user = AppUser.objects.get(uid=uid)
+    # don't send email if user is not paying or has enough balance
+    if not user.is_paying or user.balance > settings.LOW_BALANCE_EMAIL_CREDITS:
+        return
+    last_purchase = (
         AppUserTransaction.objects.filter(user=user, amount__gt=0)
         .order_by("-created_at")
         .first()
     )
-    if last_positive_transaction:
-        last_positive_transaction = last_positive_transaction.created_at
-    else:
-        last_positive_transaction = timezone.now() - datetime.timedelta(days=8)
+    email_date_cutoff = timezone.now() - datetime.timedelta(
+        days=settings.LOW_BALANCE_EMAIL_DAYS
+    )
+    # send email if user has not been sent email in last X days or last purchase was after last email sent
     if (
-        user.is_paying
-        and user.balance < settings.EMAIL_CREDITS_THRESHOLD
-        and settings.ALLOW_SENDING_CREDIT_EMAILS
-        and (
-            user.low_balance_email_sent_at == None
-            or user.low_balance_email_sent_at
-            < timezone.now() - datetime.timedelta(days=7)
-            or last_positive_transaction > user.low_balance_email_sent_at
-        )
+        # user has not been sent any email
+        not user.low_balance_email_sent_at
+        # user was sent email before X days
+        or (user.low_balance_email_sent_at < email_date_cutoff)
+        # user has made a purchase after last email sent
+        or (last_purchase and last_purchase.created_at > user.low_balance_email_sent_at)
     ):
-        total_credits_consumed = (
-            -1
-            * AppUserTransaction.objects.filter(
-                user=user,
-                created_at__gte=timezone.now() - datetime.timedelta(days=7),
-                amount__lt=0,
+        # calculate total credits consumed in last X days
+        total_credits_consumed = abs(
+            AppUserTransaction.objects.filter(
+                user=user, amount__lt=0, created_at__gte=email_date_cutoff
             ).aggregate(Sum("amount"))["amount__sum"]
+            or 0
         )
-        send_low_balance_email(
-            user=user,
-            total_credits_consumed=total_credits_consumed,
-        )
+        send_low_balance_email(user=user, total_credits_consumed=total_credits_consumed)
         user.low_balance_email_sent_at = timezone.now()
-        user.save()
+        user.save(update_fields=["low_balance_email_sent_at"])
 
 
 def send_email_on_completion(page: BasePage, sr: SavedRun):

conftest.py

@@ -9,6 +9,7 @@
 from auth import auth_backend
 from celeryapp import app
 from daras_ai_v2.base import BasePage
+from daras_ai_v2.send_email import pytest_outbox
 
 
 @pytest.fixture(scope="session")
@@ -68,6 +69,11 @@ def runner(*args, **kwargs):
             t.join()
 
 
+@pytest.fixture(autouse=True)
+def clear_pytest_outbox():
+    pytest_outbox.clear()
+
+
 # class DummyDatabaseBlocker(pytest_django.plugin._DatabaseBlocker):
 #     class _dj_db_wrapper:
 #         def ensure_connection(self):

daras_ai_v2/send_email.py

@@ -1,4 +1,5 @@
 import smtplib
+import sys
 import typing
 from email.mime.application import MIMEApplication
 from email.mime.multipart import MIMEMultipart
@@ -12,6 +13,7 @@
 from daras_ai_v2 import settings
 from daras_ai_v2.settings import templates
 from gooey_ui import UploadedFile
+from routers.billing import account_url
 
 
 def send_reported_run_email(
@@ -51,8 +53,9 @@ def send_low_balance_email(
     recipeints = "[email protected], [email protected]"
     html_body = templates.get_template("low_balance_email.html").render(
         user=user,
-        url="https://gooey.ai/account",
+        url=account_url,
         total_credits_consumed=total_credits_consumed,
+        settings=settings,
     )
     send_email_via_postmark(
         from_address=settings.SUPPORT_EMAIL,
@@ -63,6 +66,10 @@ def send_low_balance_email(
     )
 
 
+is_running_pytest = "pytest" in sys.modules
+pytest_outbox = []
+
+
 def send_email_via_postmark(
     *,
     from_address: str,
@@ -76,6 +83,21 @@ def send_email_via_postmark(
         "outbound", "gooey-ai-workflows", "announcements"
     ] = "outbound",
 ):
+    if is_running_pytest:
+        pytest_outbox.append(
+            dict(
+                from_address=from_address,
+                to_address=to_address,
+                cc=cc,
+                bcc=bcc,
+                subject=subject,
+                html_body=html_body,
+                text_body=text_body,
+                message_stream=message_stream,
+            ),
+        )
+        return
+
     r = requests.post(
         "https://api.postmarkapp.com/email",
         headers={

daras_ai_v2/settings.py

@@ -268,8 +268,9 @@
 ANON_USER_FREE_CREDITS = config("ANON_USER_FREE_CREDITS", 25, cast=int)
 LOGIN_USER_FREE_CREDITS = config("LOGIN_USER_FREE_CREDITS", 1000, cast=int)
 
-EMAIL_CREDITS_THRESHOLD = config("EMAIL_CREDITS_THRESHOLD", 500, cast=int)
-ALLOW_SENDING_CREDIT_EMAILS = config("ALLOW_SENDING_CREDIT_EMAILS", True, cast=bool)
+LOW_BALANCE_EMAIL_CREDITS = config("LOW_BALANCE_EMAIL_CREDITS", 200, cast=int)
+LOW_BALANCE_EMAIL_DAYS = config("LOW_BALANCE_EMAIL_DAYS", 7, cast=int)
+LOW_BALANCE_EMAIL_ENABLED = config("LOW_BALANCE_EMAIL_ENABLED", True, cast=bool)
 
 stripe.api_key = config("STRIPE_SECRET_KEY", None)
 STRIPE_ENDPOINT_SECRET = config("STRIPE_ENDPOINT_SECRET", None)

routers/billing.py

@@ -121,9 +121,7 @@ def account(request: Request):
         "user_credits": request.user.balance,
         "subscription": get_user_subscription(request.user),
         "is_admin": is_admin,
-        "canonical_url": str(
-            furl(settings.APP_BASE_URL) / router.url_path_for(account.__name__)
-        ),
+        "canonical_url": account_url,
     }
 
     return templates.TemplateResponse("account.html", context)

templates/low_balance_email.html

@@ -3,7 +3,8 @@
 </p>
 
 <p>
-    This is a friendly reminder that your Gooey.AI balance is now just {{ user.balance }}. Your account has consumed {{ total_credits_consumed }} credits in the last 7 days. 
+    This is a friendly reminder that your Gooey.AI balance is now just {{ user.balance }}.
+    Your account has consumed {{ total_credits_consumed }} credits in the last {{ settings.LOW_BALANCE_EMAIL_DAYS }} days.
     <br><br>
     To buy more credits, please visit <a href="{{ url }}">https://gooey.ai/account</a>.
     <br><br>

tests/test_low_balance_email_check.py

@@ -0,0 +1,127 @@
+from django.utils import timezone
+
+from app_users.models import AppUserTransaction
+from bots.models import AppUser
+from celeryapp.tasks import run_low_balance_email_check
+from daras_ai_v2 import settings
+from daras_ai_v2.send_email import pytest_outbox
+
+
+def test_dont_send_email_if_feature_is_disabled(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=True, balance=0, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = False
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_is_not_paying(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=False, balance=0, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_has_enough_balance(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=True, balance=500, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_has_been_emailed_recently(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=66,
+        is_anonymous=False,
+        low_balance_email_sent_at=timezone.now(),
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_send_email_if_user_has_been_email_recently_but_made_a_purchase(
+    transactional_db,
+):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=22,
+        is_anonymous=False,
+        low_balance_email_sent_at=timezone.now(),
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_1",
+        user=user,
+        amount=100,
+        created_at=timezone.now(),
+        end_balance=100,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_2",
+        user=user,
+        amount=-78,
+        created_at=timezone.now(),
+        end_balance=22,
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    run_low_balance_email_check(user.uid)
+
+    assert len(pytest_outbox) == 1
+    assert " 22" in pytest_outbox[0]["html_body"]
+    assert " 78" in pytest_outbox[0]["html_body"]
+
+    pytest_outbox.clear()
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_send_email(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=66,
+        is_anonymous=False,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_1",
+        user=user,
+        amount=-100,
+        created_at=timezone.now() - timezone.timedelta(days=2),
+        end_balance=150 + 66,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_2",
+        user=user,
+        amount=-150,
+        created_at=timezone.now(),
+        end_balance=66,
+    )
+
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+
+    run_low_balance_email_check(user.uid)
+    assert len(pytest_outbox) == 1
+    body = pytest_outbox[0]["html_body"]
+    assert " 66" in body
+    assert " 150" in body
+    assert " pm:unsubscribe" in body
+    assert " 100" not in body
+
+    pytest_outbox.clear()
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox