From ffe970aee535cbac0a17b45353cf1a748fe5d97a Mon Sep 17 00:00:00 2001
From: Oleksii Baranov <olebaran@redhat.com>
Date: Thu, 15 Sep 2022 13:36:58 +0200
Subject: [PATCH] Added instance_groups method to the awxkit models

Also added additional payload fields to the wfjt model.
---
 awxkit/awxkit/api/pages/schedules.py                |  8 +++++---
 .../awxkit/api/pages/workflow_job_template_nodes.py | 13 ++++++++++---
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/awxkit/awxkit/api/pages/schedules.py b/awxkit/awxkit/api/pages/schedules.py
index 34a81f878b..3ff9e1c0bb 100644
--- a/awxkit/awxkit/api/pages/schedules.py
+++ b/awxkit/awxkit/api/pages/schedules.py
@@ -48,10 +48,12 @@ class Schedules(page.PageList, Schedule):
             self.related.credentials.post(dict(id=cred.id, disassociate=True))
 
     def add_label(self, label):
-        if isinstance(label, page.Page):
-            label = label.json
         with suppress(exc.NoContent):
-            self.related.labels.post(label)
+            self.related.labels.post(dict(id=label.id))
+
+    def add_instance_group(self, instance_group):
+        with suppress(exc.NoContent):
+            self.related.instance_groups.post(dict(id=instance_group.id))
 
 
 page.register_page([resources.schedules, resources.related_schedules], Schedules)
diff --git a/awxkit/awxkit/api/pages/workflow_job_template_nodes.py b/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
index aa312cc837..3177c24e50 100644
--- a/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
+++ b/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
@@ -32,6 +32,11 @@ class WorkflowJobTemplateNode(HasCreate, base.Base):
             'extra_data',
             'identifier',
             'all_parents_must_converge',
+            # prompt fields for JTs
+            'job_slice_count',
+            'forks',
+            'timeout',
+            'execution_environment',
         )
 
         update_payload(payload, optional_fields, kwargs)
@@ -93,10 +98,12 @@ class WorkflowJobTemplateNode(HasCreate, base.Base):
         return candidates.results.pop()
 
     def add_label(self, label):
-        if isinstance(label, page.Page):
-            label = label.json
         with suppress(exc.NoContent):
-            self.related.labels.post(label)
+            self.related.labels.post(dict(id=label.id))
+
+    def add_instance_group(self, instance_group):
+        with suppress(exc.NoContent):
+            self.related.instance_groups.post(dict(id=instance_group.id))
 
 
 page.register_page(