Allow creating and updating allocations

QuanMPhm · QuanMPhm · commit 06dbb3dc78b6 · 2025-09-18T14:50:13.000-04:00
To create an allocation, a JSON payload can be uploaded to `/api/allocations`:

{
    "attributes": [
        {"attribute_type": "OpenShift Limit on CPU Quota", "value": 8},
        {"attribute_type": "OpenShift Limit on RAM Quota (MiB)", "value": 16},
    ],
    "project": {"id": project.id},
    "resources": [{"id": self.resource.id}],
    "status": "New",
}

To update an allocation, a partial JSON payload like so to `/api/allocations/{allocation id}`:

{
    "attributes": [
        {"attribute_type": "OpenShift Limit on CPU Quota", "value": 25},
        {"attribute_type": "OpenShift Limit on RAM Quota (MiB)", "value": 1600},
    ]
}
diff --git a/src/coldfront_plugin_api/serializers.py b/src/coldfront_plugin_api/serializers.py
@@ -1,14 +1,27 @@
+import logging
+
 from rest_framework import serializers
 
-from coldfront.core.allocation.models import Allocation, AllocationAttribute
-from coldfront.core.allocation.models import Project
+from coldfront.core.allocation.models import (
+    Allocation,
+    AllocationAttribute,
+    AllocationStatusChoice,
+    AllocationAttributeType,
+)
+from coldfront.core.allocation.models import Project, Resource
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
 
 class ProjectSerializer(serializers.ModelSerializer):
     class Meta:
         model = Project
         fields = ["id", "title", "pi", "description", "field_of_science", "status"]
+        read_only_fields = ["title", "pi", "description", "field_of_science", "status"]
 
+    id = serializers.IntegerField()
     pi = serializers.SerializerMethodField()
     field_of_science = serializers.SerializerMethodField()
     status = serializers.SerializerMethodField()
@@ -23,28 +36,88 @@ def get_status(self, obj: Project) -> str:
         return obj.status.name
 
 
+class AllocationAttributeSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = AllocationAttribute
+        fields = ["attribute_type", "value"]
+
+    attribute_type = (
+        serializers.SlugRelatedField(  # Peforms validation to ensure attribute exists
+            read_only=False,
+            slug_field="name",
+            queryset=AllocationAttributeType.objects.all(),
+            source="allocation_attribute_type",
+        )
+    )
+    value = serializers.CharField(read_only=False)
+
+
+class ResourceSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = Resource
+        fields = ["id", "name", "resource_type"]
+
+    id = serializers.IntegerField()
+    name = serializers.CharField(required=False)
+    resource_type = serializers.SerializerMethodField(required=False)
+
+    def get_resource_type(self, obj: Resource):
+        return obj.resource_type.name
+
+
 class AllocationSerializer(serializers.ModelSerializer):
     class Meta:
         model = Allocation
-        fields = ["id", "project", "description", "resource", "status", "attributes"]
+        fields = ["id", "project", "description", "resources", "status", "attributes"]
 
-    resource = serializers.SerializerMethodField()
+    resources = ResourceSerializer(many=True)
     project = ProjectSerializer()
-    attributes = serializers.SerializerMethodField()
-    status = serializers.SerializerMethodField()
+    attributes = AllocationAttributeSerializer(
+        many=True, source="allocationattribute_set", required=False
+    )
+    status = serializers.SlugRelatedField(
+        slug_field="name", queryset=AllocationStatusChoice.objects.all()
+    )
+
+    # TODO (Quan): What about default start/end dates? Default quantity? Description? Justification?
+
+    def create(self, validated_data):
+        project_obj = Project.objects.get(id=validated_data["project"]["id"])
+        resource_obj = Resource.objects.get(id=validated_data["resources"][0]["id"])
+        allocation = Allocation.objects.create(
+            project=project_obj, status=validated_data["status"]
+        )
+        allocation.resources.add(resource_obj)
+        allocation.save()
 
-    def get_resource(self, obj: Allocation) -> dict:
-        resource = obj.resources.first()
-        return {"name": resource.name, "resource_type": resource.resource_type.name}
+        # TODO (Quan): If the status is `Active`, do we fire the `activate_allocation`
+        # signal as well to allow creating the project on remote cluster?
 
-    def get_attributes(self, obj: Allocation):
-        attrs = AllocationAttribute.objects.filter(allocation=obj)
-        return {
-            a.allocation_attribute_type.name: obj.get_attribute(
-                a.allocation_attribute_type.name
+        for attribute in validated_data.pop("allocationattribute_set", []):
+            AllocationAttribute.objects.create(
+                allocation=allocation,
+                allocation_attribute_type=attribute["allocation_attribute_type"],
+                value=attribute["value"],
             )
-            for a in attrs
-        }
 
-    def get_status(self, obj: Allocation) -> str:
-        return obj.status.name
+        logger.info(
+            f"Created allocation {allocation.id} for project {project_obj.title}"
+        )
+        return allocation
+
+    def update(self, allocation: Allocation, validated_data):
+        """Only allow updating allocation attributes for now"""
+        # TODO (Quan) Do we want to allow updating any other allocation properties?
+        new_allocation_attributes = validated_data.pop("allocationattribute_set", [])
+        for attribute in new_allocation_attributes:
+            allocation_attribute, _ = AllocationAttribute.objects.get_or_create(
+                allocation=allocation,
+                allocation_attribute_type=attribute["allocation_attribute_type"],
+            )
+            allocation_attribute.value = attribute["value"]
+            allocation_attribute.save()
+
+        logger.info(
+            f"Updated allocation {allocation.id} for project {allocation.project.title}"
+        )
+        return allocation
diff --git a/src/coldfront_plugin_api/tests/unit/test_allocations.py b/src/coldfront_plugin_api/tests/unit/test_allocations.py
@@ -106,7 +106,7 @@ def test_filter_allocations(self):
         ).json()
         self.assertEqual(len(r_json), 3)
         self.assertEqual(
-            r_json[0]["resource"]["resource_type"], self.resource.resource_type.name
+            r_json[0]["resources"][0]["resource_type"], self.resource.resource_type.name
         )
 
         # Filter by 1 attribute with conditional or
@@ -118,8 +118,11 @@ def test_filter_allocations(self):
                 aa3.value,
             )
         ).json()
+        r_attribute_dict = {
+            a["attribute_type"]: a["value"] for a in r_json[0]["attributes"]
+        }
         self.assertEqual(len(r_json), 3)
-        self.assertIn(attributes.QUOTA_LIMITS_CPU, r_json[0]["attributes"])
+        self.assertIn(attributes.QUOTA_LIMITS_CPU, r_attribute_dict)
 
         # Filter by two allocation attributes, with conditional or
         r_json = self.admin_client.get(
@@ -130,12 +133,13 @@ def test_filter_allocations(self):
                 aa4.value,
             )
         ).json()
+        r_attribute_dict = {
+            a["attribute_type"]: a["value"] for a in r_json[0]["attributes"]
+        }
         self.assertEqual(len(r_json), 1)
+        self.assertEqual(r_attribute_dict[attributes.QUOTA_LIMITS_CPU], str(aa1.value))
         self.assertEqual(
-            r_json[0]["attributes"][attributes.QUOTA_LIMITS_CPU], aa1.value
-        )
-        self.assertEqual(
-            r_json[0]["attributes"][attributes.QUOTA_LIMITS_MEMORY], aa4.value
+            r_attribute_dict[attributes.QUOTA_LIMITS_MEMORY], str(aa4.value)
         )
 
         # Filter by non-existant attribute
@@ -146,3 +150,82 @@ def test_filter_allocations(self):
             "/api/allocations?fake_model_attribute=fake"
         ).json()
         self.assertEqual(r_json, [])
+
+    def test_create_allocation(self):
+        user = self.new_user()
+        project = self.new_project(pi=user)
+
+        payload = {
+            "attributes": [
+                {"attribute_type": "OpenShift Limit on CPU Quota", "value": 8},
+                {"attribute_type": "OpenShift Limit on RAM Quota (MiB)", "value": 16},
+            ],
+            "project": {"id": project.id},
+            "resources": [{"id": self.resource.id}],
+            "status": "New",
+        }
+
+        self.admin_client.post("/api/allocations", payload, format="json")
+
+        created_allocation = allocation_models.Allocation.objects.get(
+            project=project,
+            resources__in=[self.resource],
+        )
+        self.assertEqual(created_allocation.status.name, "New")
+
+        allocation_models.AllocationAttribute.objects.get(
+            allocation=created_allocation,
+            allocation_attribute_type=allocation_models.AllocationAttributeType.objects.get(
+                name="OpenShift Limit on CPU Quota"
+            ),
+            value=8,
+        )
+        allocation_models.AllocationAttribute.objects.get(
+            allocation=created_allocation,
+            allocation_attribute_type=allocation_models.AllocationAttributeType.objects.get(
+                name="OpenShift Limit on RAM Quota (MiB)"
+            ),
+            value=16,
+        )
+
+    def test_update_allocation(self):
+        user = self.new_user()
+        project = self.new_project(pi=user)
+        allocation = self.new_allocation(project, self.resource, 1)
+
+        # Add initial attributes
+        self.new_allocation_attribute(allocation, "OpenShift Limit on CPU Quota", 4)
+        self.new_allocation_attribute(
+            allocation, "OpenShift Limit on RAM Quota (MiB)", 8
+        )
+
+        payload = {
+            "attributes": [
+                {
+                    "attribute_type": "OpenShift Limit on CPU Quota",
+                    "value": 8,
+                },  # update CPU
+                {
+                    "attribute_type": "OpenShift Limit on RAM Quota (MiB)",
+                    "value": 16,
+                },  # update RAM
+            ],
+        }
+
+        response = self.admin_client.patch(
+            f"/api/allocations/{allocation.id}", payload, format="json"
+        )
+        self.assertEqual(response.status_code, 200)
+
+        updated_allocation = allocation_models.Allocation.objects.get(id=allocation.id)
+
+        allocation_models.AllocationAttribute.objects.get(
+            allocation=updated_allocation,
+            allocation_attribute_type__name="OpenShift Limit on CPU Quota",
+            value=8,
+        )
+        allocation_models.AllocationAttribute.objects.get(
+            allocation=updated_allocation,
+            allocation_attribute_type__name="OpenShift Limit on RAM Quota (MiB)",
+            value=16,
+        )
diff --git a/src/coldfront_plugin_api/urls.py b/src/coldfront_plugin_api/urls.py
@@ -10,7 +10,7 @@
 from coldfront_plugin_api import auth, serializers
 
 
-class AllocationViewSet(viewsets.ReadOnlyModelViewSet):
+class AllocationViewSet(viewsets.ModelViewSet):
     """
     This viewset implements the API to Coldfront's allocation object
     The API allows filtering allocations by any of Coldfront's allocation model attributes,
