From 117c7e0c42a84998c1cfa8608c403bb3665c87d4 Mon Sep 17 00:00:00 2001
From: zandercymatics <141044360+zandercymatics@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:55:50 -0600
Subject: [PATCH] Rebase logic

---
 src/registrar/models/user.py        | 13 ++++++++-----
 src/registrar/utility/csv_export.py | 13 ++-----------
 src/registrar/views/index.py        |  2 +-
 src/registrar/views/portfolios.py   |  2 +-
 4 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/registrar/models/user.py b/src/registrar/models/user.py
index 61a1a39fc..e87fb255a 100644
--- a/src/registrar/models/user.py
+++ b/src/registrar/models/user.py
@@ -4,6 +4,7 @@ from django.contrib.auth.models import AbstractUser
 from django.db import models
 from django.db.models import Q
 
+from registrar.models.domain_information import DomainInformation
 from registrar.models.user_domain_role import UserDomainRole
 
 from .domain_invitation import DomainInvitation
@@ -428,8 +429,10 @@ class User(AbstractUser):
     def is_org_user(self, request):
         has_organization_feature_flag = flag_is_active(request, "organization_feature")
         return has_organization_feature_flag and self.has_base_portfolio_permission()
-    
-    def user_domain_count(self):
-        """Returns the number of domains associated with this user through UserDomainRole"""
-        available_domains = UserDomainRole.objects.filter(user=self, domain__isnull=False)
-        return available_domains.count()
+
+    def get_user_domain_ids(self):
+        """Returns either the domains ids associated with this user on UserDomainRole or Portfolio"""
+        if self.has_base_portfolio_permission() and self.has_view_all_domains_permission():
+            return DomainInformation.objects.filter(portfolio=self.portfolio).values_list("domain_id", flat=True)
+        else:
+            return UserDomainRole.objects.filter(user=self).values_list("domain_id", flat=True)
diff --git a/src/registrar/utility/csv_export.py b/src/registrar/utility/csv_export.py
index a18d74bae..78ced6518 100644
--- a/src/registrar/utility/csv_export.py
+++ b/src/registrar/utility/csv_export.py
@@ -577,18 +577,9 @@ class DomainDataTypeUser(DomainDataType):
         if request is None or not hasattr(request, "user") or not request.user:
             # Return nothing
             return Q(id__in=[])
-
-        if (
-            request.user.has_base_portfolio_permission() and 
-            request.user.has_view_all_domains_permission()
-        ):
-            # Question: should we also include all domains in UserDomainRole as well?
-            models = DomainInformation.objects.filter(portfolio=request.user.portfolio)
         else:
-            models = UserDomainRole.objects.filter(user=request.user)
-
-        domain_ids = models.values_list("domain_id", flat=True)
-        return Q(domain__id__in=domain_ids)
+            # Get all domains the user is associated with
+            return Q(domain__id__in=request.user.get_user_domain_ids())
 
 
 class DomainDataFull(DomainExport):
diff --git a/src/registrar/views/index.py b/src/registrar/views/index.py
index 9483cdd52..0d74a81cc 100644
--- a/src/registrar/views/index.py
+++ b/src/registrar/views/index.py
@@ -8,6 +8,6 @@ def index(request):
     if request.user.is_authenticated:
         # This controls the creation of a new domain request in the wizard
         request.session["new_request"] = True
-        context["user_domain_count"] = request.user.user_domain_count()
+        context["user_domain_count"] = request.user.get_user_domain_ids().count()
 
     return render(request, "home.html", context)
diff --git a/src/registrar/views/portfolios.py b/src/registrar/views/portfolios.py
index c9f2d5ec1..7421487a0 100644
--- a/src/registrar/views/portfolios.py
+++ b/src/registrar/views/portfolios.py
@@ -29,7 +29,7 @@ class PortfolioDomainsView(PortfolioDomainsPermissionView, View):
             context["has_organization_feature_flag"] = flag_is_active(request, "organization_feature")
             portfolio = get_object_or_404(Portfolio, id=portfolio_id)
             context["portfolio"] = portfolio
-            context["user_domain_count"] = self.request.user.user_domain_count()
+            context["user_domain_count"] = self.request.user.get_user_domain_ids().count()
 
         return render(request, "portfolio_domains.html", context)