diff --git a/atst/routes/task_orders/new.py b/atst/routes/task_orders/new.py index c153364c..ee874ab4 100644 --- a/atst/routes/task_orders/new.py +++ b/atst/routes/task_orders/new.py @@ -62,6 +62,12 @@ class ShowTaskOrderWorkflow: return self._task_order + @property + def portfolio(self): + if self.task_order: + return self.task_order.portfolio + return Portfolios.get(self.user, self.portfolio_id) + @property def form(self): form_type = ( @@ -129,11 +135,6 @@ class ShowTaskOrderWorkflow: return True return False - def get_portfolio(self): - if self.task_order: - return self.task_order.portfolio - return Portfolios.get(self.user, self.portfolio_id) - class UpdateTaskOrderWorkflow(ShowTaskOrderWorkflow): def __init__( @@ -238,7 +239,7 @@ def new(screen, task_order_id=None, portfolio_id=None): } if workflow.pf_attributes_read_only(): - template_args["portfolio"] = workflow.get_portfolio() + template_args["portfolio"] = workflow.portfolio if screen == 1: workflow.form = task_order_form.AppInfoWithExistingPortfolioForm( obj=workflow.task_order diff --git a/tests/routes/task_orders/test_new_task_order.py b/tests/routes/task_orders/test_new_task_order.py index 5ed8bfc1..c99e3be4 100644 --- a/tests/routes/task_orders/test_new_task_order.py +++ b/tests/routes/task_orders/test_new_task_order.py @@ -9,6 +9,26 @@ from atst.utils.localization import translate from tests.factories import UserFactory, TaskOrderFactory, PortfolioFactory +class TestShowTaskOrderWorkflow: + def test_portfolio_when_task_order_exists(self): + portfolio = PortfolioFactory.create() + task_order = TaskOrderFactory(portfolio=portfolio) + assert portfolio.num_task_orders > 0 + + workflow = ShowTaskOrderWorkflow( + user=task_order.creator, task_order_id=task_order.id + ) + assert portfolio == workflow.portfolio + + def test_portfolio_with_portfolio_id(self): + user = UserFactory.create() + portfolio = PortfolioFactory.create(owner=user) + workflow = ShowTaskOrderWorkflow( + user=portfolio.owner, portfolio_id=portfolio.id + ) + assert portfolio == workflow.portfolio + + def test_new_task_order(client, user_session): creator = UserFactory.create() user_session() @@ -73,24 +93,6 @@ def test_to_on_pf_cannot_edit_pf_attributes(): assert second_workflow.pf_attributes_read_only() -def test_get_portfolio_when_task_order_exists(): - portfolio = PortfolioFactory.create() - task_order = TaskOrderFactory(portfolio=portfolio) - assert portfolio.num_task_orders > 0 - - workflow = ShowTaskOrderWorkflow( - user=task_order.creator, task_order_id=task_order.id - ) - assert portfolio == workflow.get_portfolio() - - -def test_get_portfolio_with_portfolio_id(): - user = UserFactory.create() - portfolio = PortfolioFactory.create(owner=user) - workflow = ShowTaskOrderWorkflow(user=portfolio.owner, portfolio_id=portfolio.id) - assert portfolio == workflow.get_portfolio() - - # TODO: this test will need to be more complicated when we add validation to # the forms def test_create_new_task_order(client, user_session, pdf_upload):