From 3182f3438e7abf1a62149382e1637e76bb2d2bbd Mon Sep 17 00:00:00 2001
From: leigh-mil <leigh-ctr@friends.dds.mil>
Date: Thu, 14 Mar 2019 12:03:42 -0400
Subject: [PATCH] Have officer invite url point to the TO view page

---
 atst/domain/task_orders.py            | 11 +++++++++++
 atst/routes/portfolios/invitations.py | 11 ++---------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/atst/domain/task_orders.py b/atst/domain/task_orders.py
index 63239679..3aa5ba88 100644
--- a/atst/domain/task_orders.py
+++ b/atst/domain/task_orders.py
@@ -185,6 +185,17 @@ class TaskOrders(object):
                 "{} is not an officer role on task orders".format(officer_type)
             )
 
+    @classmethod
+    def is_officer_for_to(cls, task_order, user):
+        if (
+            user == task_order.contracting_officer
+            or user == task_order.contracting_officer_representative
+            or user == task_order.security_officer
+        ):
+            return True
+        else:
+            return False
+
     @classmethod
     def add_dd_254(user, task_order, dd_254_data):
         dd_254 = DD254(**dd_254_data)
diff --git a/atst/routes/portfolios/invitations.py b/atst/routes/portfolios/invitations.py
index 5b1bbce6..d2d992b8 100644
--- a/atst/routes/portfolios/invitations.py
+++ b/atst/routes/portfolios/invitations.py
@@ -3,6 +3,7 @@ from flask import g, redirect, url_for, render_template
 from . import portfolios_bp
 from atst.domain.portfolios import Portfolios
 from atst.domain.invitations import Invitations
+from atst.domain.task_orders import TaskOrders
 from atst.queue import queue
 from atst.utils.flash import formatted_flash as flash
 
@@ -26,7 +27,7 @@ def accept_invitation(token):
     #   - the logged-in user has multiple roles on the TO (e.g., KO and COR)
     #   - the logged-in user has officer roles on multiple unsigned TOs
     for task_order in invite.portfolio.task_orders:
-        if g.current_user == task_order.contracting_officer:
+        if TaskOrders.is_officer_for_to(task_order, g.current_user):
             return redirect(
                 url_for(
                     "portfolios.view_task_order",
@@ -34,14 +35,6 @@ def accept_invitation(token):
                     task_order_id=task_order.id,
                 )
             )
-        elif g.current_user == task_order.contracting_officer_representative:
-            return redirect(
-                url_for("task_orders.new", screen=4, task_order_id=task_order.id)
-            )
-        elif g.current_user == task_order.security_officer:
-            return redirect(
-                url_for("task_orders.new", screen=4, task_order_id=task_order.id)
-            )
 
     return redirect(
         url_for("portfolios.show_portfolio", portfolio_id=invite.portfolio.id)