From 4ef8df016dbf6159da88fc64d7ee314613527ad2 Mon Sep 17 00:00:00 2001
From: Montana <montana-ctr@friends.dds.mil>
Date: Tue, 19 Feb 2019 11:36:59 -0500
Subject: [PATCH] Redirect user back to KO Review page if that is where they
 came from

---
 atst/routes/task_orders/new.py  | 43 ++++++++++++++++++++++++++++-----
 templates/task_orders/_new.html | 10 +++++++-
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/atst/routes/task_orders/new.py b/atst/routes/task_orders/new.py
index d9bee3b7..4c1d15bf 100644
--- a/atst/routes/task_orders/new.py
+++ b/atst/routes/task_orders/new.py
@@ -207,6 +207,33 @@ def new(screen, task_order_id=None, portfolio_id=None):
         if not TaskOrders.all_sections_complete(task_order):
             flash("task_order_draft")
 
+    if workflow.task_order:
+        ko_review_url = url_for(
+            "portfolios.ko_review",
+            portfolio_id=workflow.task_order.portfolio.id,
+            task_order_id=task_order_id,
+            _external=True,
+        )
+        redirect_url = url_for(
+            "portfolios.ko_review",
+            portfolio_id=workflow.task_order.portfolio.id,
+            task_order_id=task_order_id,
+        )
+
+        if http_request.referrer == ko_review_url:
+            return render_template(
+                workflow.template,
+                current=screen,
+                task_order_id=task_order_id,
+                task_order=workflow.task_order,
+                portfolio_id=portfolio_id,
+                screens=workflow.display_screens,
+                form=workflow.form,
+                complete=workflow.is_complete,
+                from_ko_review=True,
+                next=redirect_url,
+            )
+
     return render_template(
         workflow.template,
         current=screen,
@@ -229,15 +256,19 @@ def update(screen, task_order_id=None, portfolio_id=None):
     workflow = UpdateTaskOrderWorkflow(
         g.current_user, form_data, screen, task_order_id, portfolio_id
     )
+
     if workflow.validate():
         workflow.update()
-        return redirect(
-            url_for(
-                "task_orders.new",
-                screen=screen + 1,
-                task_order_id=workflow.task_order.id,
+        if http_request.args.get("next"):
+            return redirect(http_request.args.get("next"))
+        else:
+            return redirect(
+                url_for(
+                    "task_orders.new",
+                    screen=screen + 1,
+                    task_order_id=workflow.task_order.id,
+                )
             )
-        )
     else:
         return render_template(
             workflow.template,
diff --git a/templates/task_orders/_new.html b/templates/task_orders/_new.html
index 37f1ce74..cc25c248 100644
--- a/templates/task_orders/_new.html
+++ b/templates/task_orders/_new.html
@@ -9,6 +9,10 @@
   {% include "fragments/flash.html" %}
 
   {% block form_action %}
+    {% if from_ko_review %}
+      <form method='POST' action="{{ url_for('task_orders.update', screen=current, task_order_id=task_order_id, next=next) }}" autocomplete="off" enctype="multipart/form-data">
+    {% endif %}
+
     {% if task_order_id %}
       <form method='POST' action="{{ url_for('task_orders.new', screen=current, task_order_id=task_order_id) }}" autocomplete="off" enctype="multipart/form-data">
     {% else %}
@@ -39,7 +43,11 @@
     {% block next %}
 
       <div class='action-group'>
-        <input type='submit' class='usa-button usa-button-primary' value='Save & Continue' />
+        {% if from_ko_review %}
+          <input type='submit' class='usa-button usa-button-primary' value='Save' />
+        {% else %}
+          <input type='submit' class='usa-button usa-button-primary' value='Save & Continue' />
+        {% endif %}
       </div>
 
     {% endblock %}