From 9cf0ac03d1e2ebba289b1fdead4b6a99b1f8e72e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Szulc?= <pawel.szulc@allegro.com>
Date: Thu, 7 Nov 2024 13:39:23 +0100
Subject: [PATCH] Allow updating dchosts via dchost-list endpoint

---
 src/ralph/assets/api/views.py      | 26 +++++++++++++++++++++++---
 src/ralph/assets/tests/test_api.py | 25 +++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/ralph/assets/api/views.py b/src/ralph/assets/api/views.py
index a1109c5df7..213a5a4886 100644
--- a/src/ralph/assets/api/views.py
+++ b/src/ralph/assets/api/views.py
@@ -2,6 +2,7 @@
 import django_filters
 from django.db.models import Prefetch
 from rest_framework.exceptions import ValidationError
+from rest_framework.permissions import SAFE_METHODS
 
 from ralph.api import RalphAPIViewSet
 from ralph.api.filters import BooleanFilter
@@ -10,9 +11,11 @@
 from ralph.assets.api import serializers
 from ralph.assets.api.filters import NetworkableObjectFilters
 from ralph.assets.models import BaseObject
+from ralph.data_center.models import Cluster, DataCenterAsset
 from ralph.licences.api import BaseObjectLicenceViewSet
 from ralph.licences.models import BaseObjectLicence
 from ralph.networks.models import IPAddress
+from ralph.virtual.models import CloudHost, VirtualServer
 
 
 class BusinessSegmentViewSet(RalphAPIViewSet):
@@ -255,10 +258,10 @@ class Meta(NetworkableObjectFilters.Meta):
 # TODO: move to data_center and use DCHost proxy model
 class DCHostViewSet(BaseObjectViewSetMixin, RalphAPIViewSet):
     queryset = (
-        BaseObject.polymorphic_objects
+        BaseObject.polymorphic_objects.dc_hosts()
     )
     serializer_class = serializers.DCHostSerializer
-    http_method_names = ["get", "options", "head"]
+    http_method_names = ["get", "options", "head", "patch", "post"]
     filter_fields = [
         "id",
         "service_env",
@@ -296,9 +299,26 @@ class DCHostViewSet(BaseObjectViewSetMixin, RalphAPIViewSet):
     }
     additional_filter_class = DCHostFilterSet
 
+    def get_serializer_class(self, *args, **kwargs):
+        if self.request.method not in SAFE_METHODS:
+            obj_ = self.get_object()
+            if isinstance(obj_, VirtualServer):
+                from ralph.virtual.api import VirtualServerSaveSerializer
+                return VirtualServerSaveSerializer
+            elif isinstance(obj_, DataCenterAsset):
+                from ralph.data_center.api.serializers import DataCenterAssetSaveSerializer
+                return DataCenterAssetSaveSerializer
+            elif isinstance(obj_, CloudHost):
+                from ralph.virtual.api import SaveCloudHostSerializer
+                return SaveCloudHostSerializer
+            elif isinstance(obj_, Cluster):
+                from ralph.data_center.api.serializers import ClusterSerializer
+                return ClusterSerializer
+        return serializers.DCHostSerializer
+
     def get_queryset(self):
         return (
-            self.queryset.dc_hosts()
+            self.queryset
             .select_related(*self.select_related)
             .polymorphic_select_related(Cluster=['type'])
             .polymorphic_prefetch_related(
diff --git a/src/ralph/assets/tests/test_api.py b/src/ralph/assets/tests/test_api.py
index cb3cd229fd..b30b507be1 100644
--- a/src/ralph/assets/tests/test_api.py
+++ b/src/ralph/assets/tests/test_api.py
@@ -1089,6 +1089,31 @@ def test_filter_by_env_name(self):
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.data['count'], 1)
 
+    def test_patch_dchost_virtual_server(self):
+        new_hypervisor = DataCenterAssetFullFactory()
+        url = reverse('dchost-detail', args=(self.virtual.id,))
+        data = {
+            'hostname': 'new-hostname',
+            'hypervisor': new_hypervisor.id,
+        }
+        response = self.client.patch(url, data, format='json')
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.virtual.refresh_from_db()
+        self.assertEqual(self.virtual.hostname, 'new-hostname')
+        self.assertEqual(self.virtual.parent.id, new_hypervisor.id)
+
+    def test_patch_dchost_cloudhost(self):
+        new_hypervisor = DataCenterAssetFullFactory()
+        url = reverse('dchost-detail', args=(self.cloud_host.id,))
+        data = {
+            'hostname': 'new-hostname',
+            'hypervisor': new_hypervisor.id,
+        }
+        response = self.client.patch(url, data, format='json')
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.cloud_host.refresh_from_db()
+        self.assertEqual(self.cloud_host.hostname, 'new-hostname')
+        self.assertEqual(self.cloud_host.hypervisor.id, new_hypervisor.id)
 
 class ConfigurationModuleAPITests(RalphAPITestCase):
     def setUp(self):