diff --git a/temba/flows/tests.py b/temba/flows/tests.py index 071057d30a8..4edaa2dd8f9 100644 --- a/temba/flows/tests.py +++ b/temba/flows/tests.py @@ -21,7 +21,7 @@ from temba.classifiers.models import Classifier from temba.contacts.models import URN, Contact, ContactField, ContactGroup, ContactURN from temba.globals.models import Global -from temba.msgs.models import SystemLabel, SystemLabelCount +from temba.msgs.models import SystemLabel from temba.orgs.integrations.dtone import DTOneType from temba.orgs.models import Export from temba.templates.models import TemplateTranslation @@ -1933,7 +1933,7 @@ def test_preview_start(self, mr_mocks): ) # if we have too many messages in our outbox we should block - SystemLabelCount.objects.create(org=self.org, label_type=SystemLabel.TYPE_OUTBOX, count=1_000_001) + self.org.counts.create(scope=f"msgs:folder:{SystemLabel.TYPE_OUTBOX}", count=1_000_001) preview_url = reverse("flows.flow_preview_start", args=[flow.id]) mr_mocks.flow_start_preview(query="age > 30", total=1000) @@ -1950,7 +1950,7 @@ def test_preview_start(self, mr_mocks): ], response.json()["blockers"], ) - self.org.system_labels.all().delete() + self.org.counts.prefix("msgs:folder:").delete() # check warning for lots of contacts preview_url = reverse("flows.flow_preview_start", args=[flow.id]) diff --git a/temba/msgs/models.py b/temba/msgs/models.py index 1f9e6cb7f8f..f08c77198bd 100644 --- a/temba/msgs/models.py +++ b/temba/msgs/models.py @@ -815,12 +815,9 @@ class SystemLabel: ) @classmethod - def get_counts(cls, org, use_new=False): - if use_new: - counts = org.counts.prefix("msgs:folder:").scope_totals() - return {lb: counts.get(f"msgs:folder:{lb}", 0) for lb, n in cls.TYPE_CHOICES} - - return SystemLabelCount.get_totals(org) + def get_counts(cls, org): + counts = org.counts.prefix("msgs:folder:").scope_totals() + return {lb: counts.get(f"msgs:folder:{lb}", 0) for lb, n in cls.TYPE_CHOICES} @classmethod def get_queryset(cls, org, label_type): @@ -897,7 +894,7 @@ def get_archive_query(cls, label_type: str) -> dict: class SystemLabelCount(BaseSquashableCount): """ - Counts of messages/broadcasts/calls maintained by database level triggers + TODO drop """ squash_over = ("org_id", "label_type") @@ -905,17 +902,6 @@ class SystemLabelCount(BaseSquashableCount): org = models.ForeignKey(Org, on_delete=models.PROTECT, related_name="system_labels") label_type = models.CharField(max_length=1, choices=SystemLabel.TYPE_CHOICES) - @classmethod - def get_totals(cls, org): - """ - Gets all system label counts by type for the given org - """ - counts = cls.objects.filter(org=org).values_list("label_type").annotate(count_sum=Sum("count")) - counts_by_type = {c[0]: c[1] for c in counts} - - # for convenience, include all label types - return {lb: counts_by_type.get(lb, 0) for lb, n in SystemLabel.TYPE_CHOICES} - class Meta: indexes = [models.Index(fields=("org", "label_type", "is_squashed"))] diff --git a/temba/msgs/tests.py b/temba/msgs/tests.py index 49590984d04..bf34bef7233 100644 --- a/temba/msgs/tests.py +++ b/temba/msgs/tests.py @@ -14,20 +14,10 @@ from temba.archives.models import Archive from temba.channels.models import ChannelCount, ChannelLog from temba.flows.models import Flow -from temba.msgs.models import ( - Attachment, - Broadcast, - Label, - LabelCount, - Media, - MessageExport, - Msg, - OptIn, - SystemLabel, - SystemLabelCount, -) +from temba.msgs.models import Attachment, Broadcast, Label, LabelCount, Media, MessageExport, Msg, OptIn, SystemLabel from temba.msgs.views import ScheduleForm from temba.orgs.models import Export +from temba.orgs.tasks import squash_item_counts from temba.schedules.models import Schedule from temba.templates.models import TemplateTranslation from temba.tests import CRUDLTestMixin, TembaTest, mock_mailroom, mock_uuids @@ -2466,7 +2456,7 @@ def test_preview(self, mr_mocks): # if we have too many messages in our outbox we should block mr_mocks.msg_broadcast_preview(query="age > 30", total=2) - SystemLabelCount.objects.create(org=self.org, label_type=SystemLabel.TYPE_OUTBOX, count=1_000_001) + self.org.counts.create(scope=f"msgs:folder:{SystemLabel.TYPE_OUTBOX}", count=1_000_001) response = self.client.post(preview_url, {"query": "age > 30"}, content_type="application/json") self.assertEqual( [ @@ -2474,7 +2464,7 @@ def test_preview(self, mr_mocks): ], response.json()["blockers"], ) - self.org.system_labels.all().delete() + self.org.counts.prefix("msgs:folder:").delete() # if we release our send channel we can't send a broadcast self.channel.release(self.admin) @@ -2899,8 +2889,7 @@ def test_get_archive_query(self): def test_get_counts(self): def assert_counts(org, expected: dict): - self.assertEqual(SystemLabel.get_counts(org, use_new=False), expected) - self.assertEqual(SystemLabel.get_counts(org, use_new=True), expected) + self.assertEqual(SystemLabel.get_counts(org), expected) assert_counts( self.org, @@ -3021,10 +3010,10 @@ def assert_counts(org, expected: dict): }, ) - self.assertEqual(SystemLabelCount.objects.all().count(), 25) + self.assertEqual(self.org.counts.count(), 25) # squash our counts - squash_msg_counts() + squash_item_counts() assert_counts( self.org, @@ -3040,8 +3029,8 @@ def assert_counts(org, expected: dict): }, ) - # we should only have one system label per type with no-zero count - self.assertEqual(SystemLabelCount.objects.all().count(), 5) + # we should only have one count per folder with non-zero count + self.assertEqual(self.org.counts.count(), 5) class TagsTest(TembaTest): @@ -3188,7 +3177,7 @@ def setUpBeforeMigration(self, apps): def test_migration(self): def assert_counts(org, expected: dict): - self.assertEqual(SystemLabel.get_counts(org, use_new=True), expected) + self.assertEqual(SystemLabel.get_counts(org), expected) assert_counts( self.org,