From 59249d013be6ced77aed57891c8c884f5cc63d6c Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Sat, 5 Aug 2017 22:38:35 -0400
Subject: [PATCH] use standard type string algorithm in scheduler

---
 awx/main/scheduler/__init__.py           | 21 +++------------------
 awx/main/tests/unit/utils/test_common.py | 22 ++++++++++++++++++++++
 2 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/awx/main/scheduler/__init__.py b/awx/main/scheduler/__init__.py
index 8f54654727..1af72ff4c1 100644
--- a/awx/main/scheduler/__init__.py
+++ b/awx/main/scheduler/__init__.py
@@ -18,6 +18,7 @@ from awx.main.models import * # noqa
 #from awx.main.scheduler.dag_simple import SimpleDAG
 from awx.main.scheduler.dag_workflow import WorkflowDAG
 from awx.main.utils.pglock import advisory_lock
+from awx.main.utils import get_type_for_model
 from awx.main.signals import disable_activity_stream
 
 from awx.main.scheduler.dependency_graph import DependencyGraph
@@ -64,22 +65,6 @@ class TaskManager():
                            key=lambda task: task.created)
         return all_tasks
 
-    @classmethod
-    def get_node_type(cls, obj):
-        if type(obj) == Job:
-            return "job"
-        elif type(obj) == AdHocCommand:
-            return "ad_hoc_command"
-        elif type(obj) == InventoryUpdate:
-            return "inventory_update"
-        elif type(obj) == ProjectUpdate:
-            return "project_update"
-        elif type(obj) == SystemJob:
-            return "system_job"
-        elif type(obj) == WorkflowJob:
-            return "workflow_job"
-        return "unknown"
-
     '''
     Tasks that are running and SHOULD have a celery task.
     '''
@@ -190,10 +175,10 @@ class TaskManager():
         from awx.main.tasks import handle_work_error, handle_work_success
 
         task_actual = {
-            'type':self.get_node_type(task),
+            'type': get_type_for_model(type(task)),
             'id': task.id,
         }
-        dependencies = [{'type': self.get_node_type(t), 'id': t.id} for t in dependent_tasks]
+        dependencies = [{'type': get_type_for_model(type(t)), 'id': t.id} for t in dependent_tasks]
 
         error_handler = handle_work_error.s(subtasks=[task_actual] + dependencies)
         success_handler = handle_work_success.s(task_actual=task_actual)
diff --git a/awx/main/tests/unit/utils/test_common.py b/awx/main/tests/unit/utils/test_common.py
index f8592d844b..44e0461a9a 100644
--- a/awx/main/tests/unit/utils/test_common.py
+++ b/awx/main/tests/unit/utils/test_common.py
@@ -8,6 +8,15 @@ from uuid import uuid4
 
 from awx.main.utils import common
 
+from awx.main.models import (
+    Job,
+    AdHocCommand,
+    InventoryUpdate,
+    ProjectUpdate,
+    SystemJob,
+    WorkflowJob
+)
+
 
 @pytest.mark.parametrize('input_, output', [
     ({"foo": "bar"}, {"foo": "bar"}),
@@ -27,3 +36,16 @@ def test_set_environ():
         assert set(os.environ.keys()) - set(old_environ.keys()) == set([key])
     assert os.environ == old_environ
     assert key not in os.environ
+
+
+# Cases relied on for scheduler dependent jobs list
+@pytest.mark.parametrize('model,name', [
+    (Job, 'job'),
+    (AdHocCommand, 'ad_hoc_command'),
+    (InventoryUpdate, 'inventory_update'),
+    (ProjectUpdate, 'project_update'),
+    (SystemJob, 'system_job'),
+    (WorkflowJob, 'workflow_job')
+])
+def test_get_type_for_model(model, name):
+    assert common.get_type_for_model(model) == name