From 9019792bc8363f64566604270db11386cdc3bf93 Mon Sep 17 00:00:00 2001
From: zandercymatics <141044360+zandercymatics@users.noreply.github.com>
Date: Fri, 6 Sep 2024 15:31:26 -0600
Subject: [PATCH] Fix existing tests

---
 src/registrar/registrar_middleware.py       | 42 ++++++++++-----------
 src/registrar/tests/test_models.py          |  2 +-
 src/registrar/tests/test_views_portfolio.py |  2 +-
 3 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/src/registrar/registrar_middleware.py b/src/registrar/registrar_middleware.py
index 5323d513f..bdc76bd25 100644
--- a/src/registrar/registrar_middleware.py
+++ b/src/registrar/registrar_middleware.py
@@ -6,7 +6,7 @@ import logging
 from urllib.parse import parse_qs
 from django.urls import reverse
 from django.http import HttpResponseRedirect
-from registrar.models.user import User
+from registrar.models import User, Portfolio
 from waffle.decorators import flag_is_active
 
 from registrar.models.utility.generic_helper import replace_url_queryparams
@@ -144,28 +144,17 @@ class CheckPortfolioMiddleware:
         if not request.user.is_authenticated:
             return None
 
-        old_updated_at = None
-        if request.session.get("portfolio"):
-            old_updated_at = request.session.get("portfolio__updated_at")
-            request.session["portfolio__updated_at"] = request.session.get("portfolio").updated_at
+        portfolio = request.session.get("portfolio")
 
-        should_update_portfolio = (
-            not request.session.get("portfolio") or old_updated_at != request.session.get("portfolio__updated_at")
-        )
-        if request.user.is_org_user(request) or should_update_portfolio:
-            # if multiple portfolios are allowed for this user
-            if flag_is_active(request, "multiple_portfolios"):
-                # NOTE: we will want to change later to have a workflow for selecting
-                # portfolio and another for switching portfolio; for now, select first
-                request.session["portfolio"] = request.user.get_first_portfolio()
-            elif flag_is_active(request, "organization_feature"):
-                request.session["portfolio"] = request.user.get_first_portfolio()
-            else:
-                request.session["portfolio"] = None
-        else: 
-            # Edge case: waffle flag is changed while the user is logged in
-            if not request.user.is_org_user(request) and request.session.get("portfolio"):
-                request.session["portfolio"] = None
+        # if multiple portfolios are allowed for this user
+        if flag_is_active(request, "organization_feature"):
+            old_updated_at = request.session.get("portfolio__updated_at")
+            request.session["portfolio__updated_at"] = portfolio.updated_at if portfolio else None
+            if request.user.is_org_user(request) or old_updated_at != request.session.get("portfolio__updated_at"):
+                self.set_portfolio_in_session(request)
+        elif request.session.get("portfolio"):
+            # Edge case: User disables flag while already logged in
+            request.session["portfolio"] = None
 
         if request.session.get("portfolio"):
             if current_path == self.home:
@@ -173,7 +162,14 @@ class CheckPortfolioMiddleware:
                     portfolio_redirect = reverse("domains")
                 else:
                     portfolio_redirect = reverse("no-portfolio-domains")
-
                 return HttpResponseRedirect(portfolio_redirect)
 
         return None
+
+    def set_portfolio_in_session(self, request):
+        # NOTE: we will want to change later to have a workflow for selecting
+        # portfolio and another for switching portfolio; for now, select first
+        if flag_is_active(request, "multiple_portfolios"):
+            request.session["portfolio"] = request.user.get_first_portfolio()
+        else:
+            request.session["portfolio"] = request.user.get_first_portfolio()
diff --git a/src/registrar/tests/test_models.py b/src/registrar/tests/test_models.py
index 936e644d0..b8cf2a824 100644
--- a/src/registrar/tests/test_models.py
+++ b/src/registrar/tests/test_models.py
@@ -1135,7 +1135,7 @@ class TestPortfolioInvitations(TestCase):
         self.portfolio, _ = Portfolio.objects.get_or_create(creator=self.user2, organization_name="Hotel California")
         self.portfolio_role_base = UserPortfolioRoleChoices.ORGANIZATION_MEMBER
         self.portfolio_role_admin = UserPortfolioRoleChoices.ORGANIZATION_ADMIN
-        self.portfolio_permission_1 = UserPortfolioPermissionChoices.VIEW_CREATED_REQUESTS
+        self.portfolio_permission_1 = UserPortfolioPermissionChoices.VIEW_ALL_REQUESTS
         self.portfolio_permission_2 = UserPortfolioPermissionChoices.EDIT_REQUESTS
         self.invitation, _ = PortfolioInvitation.objects.get_or_create(
             email=self.email,
diff --git a/src/registrar/tests/test_views_portfolio.py b/src/registrar/tests/test_views_portfolio.py
index 7d3c6f3be..10c8724f5 100644
--- a/src/registrar/tests/test_views_portfolio.py
+++ b/src/registrar/tests/test_views_portfolio.py
@@ -78,7 +78,7 @@ class TestPortfolio(WebTest):
     def test_middleware_does_not_redirect_if_no_permission(self):
         """Test that user with no portfolio permission is not redirected when attempting to access home"""
         self.app.set_user(self.user.username)
-        portfolio_permission, _ = UserPortfolioPermission.objects.get_or_create(
+        UserPortfolioPermission.objects.get_or_create(
             user=self.user, portfolio=self.portfolio, additional_permissions=[]
         )
         self.user.portfolio = self.portfolio