diff --git a/app/models.py b/app/models.py index f3a6cabc..cc8f7015 100644 --- a/app/models.py +++ b/app/models.py @@ -208,6 +208,21 @@ class User(db.Model, ModelMixin, UserMixin): self.promo_codes = ",".join(current_promo_codes) + def suggested_emails(self) -> (str, [str]): + """return suggested email and other email choices """ + all_gen_emails = [ge.email for ge in GenEmail.filter_by(user_id=self.id)] + if self.can_create_new_email(): + # create a new email + suggested_gen_email = generate_email() + else: + # pick an email from the list of gen emails + suggested_gen_email = random.choice(all_gen_emails) + + return ( + suggested_gen_email, + list(set(all_gen_emails).difference(set([suggested_gen_email]))), + ) + class ActivationCode(db.Model, ModelMixin): """For activate user account""" diff --git a/tests/test_models.py b/tests/test_models.py index bcc21df9..e2221757 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -15,3 +15,42 @@ def test_profile_picture_url(flask_client): db.session.commit() assert user.profile_picture_url() == "http://sl.test/static/default-avatar.png" + + +def test_suggested_emails_for_user_who_can_create_new_email(flask_client): + user = User.create( + email="a@b.c", password="password", name="Test User", activated=True + ) + db.session.commit() + + suggested_email, other_emails = user.suggested_emails() + assert suggested_email + assert len(other_emails) == 1 + + # the suggested email is new and not exist in GenEmail + assert GenEmail.get_by(email=suggested_email) is None + + # all other emails are generated emails + assert GenEmail.get_by(email=other_emails[0]) + + +def test_suggested_emails_for_user_who_cannot_create_new_email(flask_client): + user = User.create( + email="a@b.c", password="password", name="Test User", activated=True + ) + user.plan = PlanEnum.free + db.session.commit() + + # make sure user runs out of quota to create new email + for i in range(MAX_NB_EMAIL_FREE_PLAN): + GenEmail.create_new_gen_email(user_id=user.id) + db.session.commit() + + suggested_email, other_emails = user.suggested_emails() + + # the suggested email is chosen from existing GenEmail + assert GenEmail.get_by(email=suggested_email) + + # all other emails are generated emails + for email in other_emails: + assert GenEmail.get_by(email=email)