From d5afbc20ac903ce0887f65b3570d4eeb56886484 Mon Sep 17 00:00:00 2001
From: Matthew Jones <mat@matburt.net>
Date: Tue, 10 May 2016 14:48:55 -0400
Subject: [PATCH] Require cloud on inventory source for notification

Inventory Sources must now be a cloud type in order to associate a
notification template with it.
---
 awx/api/views.py                            | 18 ++++++++++--------
 awx/main/tests/functional/test_inventory.py | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 8 deletions(-)
 create mode 100644 awx/main/tests/functional/test_inventory.py

diff --git a/awx/api/views.py b/awx/api/views.py
index dec0adfd75..6674794d12 100644
--- a/awx/api/views.py
+++ b/awx/api/views.py
@@ -1998,18 +1998,20 @@ class InventorySourceNotificationTemplatesAnyList(SubListCreateAttachDetachAPIVi
     parent_model = InventorySource
     relationship = 'notification_templates_any'
 
-class InventorySourceNotificationTemplatesErrorList(SubListCreateAttachDetachAPIView):
+    def post(self, request, *args, **kwargs):
+        parent = self.get_parent_object()
+        if parent.source not in CLOUD_INVENTORY_SOURCES:
+            return Response(dict(msg="Notification Templates can only be assigned when source is one of {}"
+                                 .format(CLOUD_INVENTORY_SOURCES, parent.source)),
+                            status=status.HTTP_400_BAD_REQUEST)
+        return super(InventorySourceNotificationTemplatesAnyList, self).post(request, *args, **kwargs)
+
+class InventorySourceNotificationTemplatesErrorList(InventorySourceNotificationTemplatesAnyList):
 
-    model = NotificationTemplate
-    serializer_class = NotificationTemplateSerializer
-    parent_model = InventorySource
     relationship = 'notification_templates_error'
 
-class InventorySourceNotificationTemplatesSuccessList(SubListCreateAttachDetachAPIView):
+class InventorySourceNotificationTemplatesSuccessList(InventorySourceNotificationTemplatesAnyList):
 
-    model = NotificationTemplate
-    serializer_class = NotificationTemplateSerializer
-    parent_model = InventorySource
     relationship = 'notification_templates_success'
 
 class InventorySourceHostsList(SubListAPIView):
diff --git a/awx/main/tests/functional/test_inventory.py b/awx/main/tests/functional/test_inventory.py
new file mode 100644
index 0000000000..768162b103
--- /dev/null
+++ b/awx/main/tests/functional/test_inventory.py
@@ -0,0 +1,19 @@
+import pytest
+
+from django.core.urlresolvers import reverse
+
+@pytest.mark.django_db
+def test_inventory_source_notification_on_cloud_only(get, post, group, user, notification_template):
+    u = user('admin', True)
+    g_cloud = group('cloud')
+    g_not = group('not_cloud')
+    cloud_is = g_cloud.inventory_source
+    not_is = g_not.inventory_source
+    cloud_is.source = 'ec2'
+    cloud_is.save()
+    url = reverse('api:inventory_source_notification_templates_any_list', args=(cloud_is.id,))
+    response = post(url, dict(id=notification_template.id), u)
+    assert response.status_code == 204
+    url = reverse('api:inventory_source_notification_templates_success_list', args=(not_is.id,))
+    response = post(url, dict(id=notification_template.id), u)
+    assert response.status_code == 400