From 8acbb85cfed02f3dec98490a1ff00d0f1b411f65 Mon Sep 17 00:00:00 2001
From: katrinan029 <katrinan@usc.edu>
Date: Wed, 25 Sep 2024 16:08:52 +0000
Subject: [PATCH] feat: update enterprise group model and serializer

---
 enterprise/api/v1/serializers.py              | 18 +++++++++++++--
 enterprise/constants.py                       |  6 +++++
 ...221_enterprisegroup_group_type_and_more.py | 23 +++++++++++++++++++
 enterprise/models.py                          | 11 +++++++++
 tests/test_enterprise/api/test_views.py       |  2 ++
 5 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 enterprise/migrations/0221_enterprisegroup_group_type_and_more.py

diff --git a/enterprise/api/v1/serializers.py b/enterprise/api/v1/serializers.py
index 36ffed664..8e04c6a9c 100644
--- a/enterprise/api/v1/serializers.py
+++ b/enterprise/api/v1/serializers.py
@@ -23,7 +23,12 @@
 from enterprise import models, utils  # pylint: disable=cyclic-import
 from enterprise.api.v1.fields import Base64EmailCSVField
 from enterprise.api_client.lms import ThirdPartyAuthApiClient
-from enterprise.constants import ENTERPRISE_ADMIN_ROLE, ENTERPRISE_PERMISSION_GROUPS, DefaultColors
+from enterprise.constants import (
+    ENTERPRISE_ADMIN_ROLE,
+    ENTERPRISE_PERMISSION_GROUPS,
+    DefaultColors,
+    GROUP_MEMBERSHIP_ACCEPTED_STATUS,
+)
 from enterprise.logging import getEnterpriseLogger
 from enterprise.models import (
     AdminNotification,
@@ -632,7 +637,16 @@ class EnterpriseGroupSerializer(serializers.ModelSerializer):
     """
     class Meta:
         model = models.EnterpriseGroup
-        fields = ('enterprise_customer', 'name', 'uuid', 'applies_to_all_contexts')
+        fields = (
+            'enterprise_customer', 'name', 'uuid', 'applies_to_all_contexts',
+            'accepted_members_count', 'group_type')
+
+    accepted_members_count = serializers.SerializerMethodField()
+
+    def get_accepted_members_count(self, obj):
+        "Returns count for accepted members"
+        all_members = obj.get_all_learners().filter(status=GROUP_MEMBERSHIP_ACCEPTED_STATUS)
+        return len(all_members)
 
 
 class EnterpriseGroupMembershipSerializer(serializers.ModelSerializer):
diff --git a/enterprise/constants.py b/enterprise/constants.py
index d38e88dce..87d2cec4a 100644
--- a/enterprise/constants.py
+++ b/enterprise/constants.py
@@ -263,6 +263,12 @@ class FulfillmentTypes:
     (GROUP_MEMBERSHIP_INTERNAL_API_ERROR_STATUS, 'Internal API error'),
     (GROUP_MEMBERSHIP_EMAIL_ERROR_STATUS, 'Email error')
 )
+GROUP_TYPE_BUDGET = 'budget'
+GROUP_TYPE_FLEX = 'flex'
+GROUP_TYPE_CHOICES = (
+    (GROUP_TYPE_BUDGET, 'Budget'),
+    (GROUP_TYPE_FLEX, 'Flex')
+)
 
 ENTITY_ID_REGEX = r"<(\w+:)?EntityDescriptor.*?entityID=['\"](.*?)['\"].*?>"
 
diff --git a/enterprise/migrations/0221_enterprisegroup_group_type_and_more.py b/enterprise/migrations/0221_enterprisegroup_group_type_and_more.py
new file mode 100644
index 000000000..090cb9da0
--- /dev/null
+++ b/enterprise/migrations/0221_enterprisegroup_group_type_and_more.py
@@ -0,0 +1,23 @@
+# Generated by Django 4.2.15 on 2024-09-24 21:39
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('enterprise', '0220_alter_updateroleassignmentswithcustomersconfig_role'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='enterprisegroup',
+            name='group_type',
+            field=models.CharField(blank=True, choices=[('budget', 'Budget'), ('flex', 'Flex')], default='flex', help_text='The type of enterprise group', max_length=20, null=True, verbose_name='Group Type'),
+        ),
+        migrations.AddField(
+            model_name='historicalenterprisegroup',
+            name='group_type',
+            field=models.CharField(blank=True, choices=[('budget', 'Budget'), ('flex', 'Flex')], default='flex', help_text='The type of enterprise group', max_length=20, null=True, verbose_name='Group Type'),
+        ),
+    ]
diff --git a/enterprise/models.py b/enterprise/models.py
index 74d3b03a8..851b9b754 100644
--- a/enterprise/models.py
+++ b/enterprise/models.py
@@ -54,6 +54,8 @@
     GROUP_MEMBERSHIP_ACCEPTED_STATUS,
     GROUP_MEMBERSHIP_PENDING_STATUS,
     GROUP_MEMBERSHIP_STATUS_CHOICES,
+    GROUP_TYPE_CHOICES,
+    GROUP_TYPE_FLEX,
     MAX_INVITE_KEYS,
     DefaultColors,
     FulfillmentTypes,
@@ -4405,6 +4407,15 @@ class EnterpriseGroup(TimeStampedModel, SoftDeletableModel):
             "When enabled, all learners connected to the org will be considered a member."
         )
     )
+    group_type = models.CharField(
+        verbose_name="Group Type",
+        max_length=20,
+        blank=True,
+        null=True,
+        choices=GROUP_TYPE_CHOICES,
+        default=GROUP_TYPE_FLEX,
+        help_text=_("The type of enterprise group"),
+    )
 
     history = HistoricalRecords()
 
diff --git a/tests/test_enterprise/api/test_views.py b/tests/test_enterprise/api/test_views.py
index aa45b85f7..73f295858 100644
--- a/tests/test_enterprise/api/test_views.py
+++ b/tests/test_enterprise/api/test_views.py
@@ -8042,6 +8042,8 @@ def test_successful_list_groups(self):
         )
         response = self.client.get(url)
         assert response.json().get('count') == 2
+        assert response.json().get('results')[0].get('group_type') == 'flex'
+        assert response.json().get('results')[0].get('accepted_members_count') == 10
 
     def test_successful_retrieve_group(self):
         """