From 2589be3b6118e5ebfe40f3fb20b0042226ef760d Mon Sep 17 00:00:00 2001 From: Dev Aggarwal Date: Wed, 28 Feb 2024 20:45:40 +0530 Subject: [PATCH] write tests for low balance email check, add indexes, reduce cuttoff to 200 credits avoid hardcoding urls and dates --- app_users/models.py | 4 + celeryapp/tasks.py | 65 +++++++------ conftest.py | 6 ++ daras_ai_v2/send_email.py | 24 ++++- daras_ai_v2/settings.py | 5 +- routers/billing.py | 4 +- templates/low_balance_email.html | 3 +- tests/test_low_balance_email_check.py | 127 ++++++++++++++++++++++++++ 8 files changed, 198 insertions(+), 40 deletions(-) create mode 100644 tests/test_low_balance_email_check.py diff --git a/app_users/models.py b/app_users/models.py index d4f714e96..4c7e716e2 100644 --- a/app_users/models.py +++ b/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})" diff --git a/celeryapp/tasks.py b/celeryapp/tasks.py index 777aa27ec..2cf56fcdc 100644 --- a/celeryapp/tasks.py +++ b/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): diff --git a/conftest.py b/conftest.py index dba333885..d535b2b47 100644 --- a/conftest.py +++ b/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): diff --git a/daras_ai_v2/send_email.py b/daras_ai_v2/send_email.py index fc2445f41..aaf1eb59b 100644 --- a/daras_ai_v2/send_email.py +++ b/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 = "support@gooey.ai, devs@gooey.ai" 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={ diff --git a/daras_ai_v2/settings.py b/daras_ai_v2/settings.py index 42ad40e16..8aed12e7b 100644 --- a/daras_ai_v2/settings.py +++ b/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) diff --git a/routers/billing.py b/routers/billing.py index bf660325b..39dfa61fb 100644 --- a/routers/billing.py +++ b/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) diff --git a/templates/low_balance_email.html b/templates/low_balance_email.html index d3215c861..ca4a65e61 100644 --- a/templates/low_balance_email.html +++ b/templates/low_balance_email.html @@ -3,7 +3,8 @@

- 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.

To buy more credits, please visit https://gooey.ai/account.

diff --git a/tests/test_low_balance_email_check.py b/tests/test_low_balance_email_check.py new file mode 100644 index 000000000..66203a19b --- /dev/null +++ b/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