From c61cc3fc47a41d5986a5bc89a640f312a502f12e Mon Sep 17 00:00:00 2001
From: erikvw <ew2789@gmail.com>
Date: Mon, 13 Jan 2025 18:44:10 -0600
Subject: [PATCH] slight refactor for

---
 edc_sites/model_mixins/site_model_mixin.py    | 46 ++++++++++---------
 .../utils/valid_site_for_subject_or_raise.py  | 25 ++++++----
 2 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/edc_sites/model_mixins/site_model_mixin.py b/edc_sites/model_mixins/site_model_mixin.py
index 580f635..10b180a 100644
--- a/edc_sites/model_mixins/site_model_mixin.py
+++ b/edc_sites/model_mixins/site_model_mixin.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from django.conf import settings
+from django.contrib.sites.models import Site
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import models, transaction
 
@@ -10,9 +9,6 @@
 from ..site import sites
 from ..utils import get_site_model_cls
 
-if TYPE_CHECKING:
-    from django.contrib.sites.models import Site
-
 
 class SiteModelMixinError(Exception):
     pass
@@ -23,7 +19,6 @@ class SiteModelMixin(models.Model):
         "sites.site",
         on_delete=models.PROTECT,
         null=True,
-        # editable=False,
         related_name="+",
     )
 
@@ -33,36 +28,43 @@ def save(self, *args, **kwargs):
         self.update_site_on_save(*args, **kwargs)
         super().save(*args, **kwargs)
 
-    def update_site_on_save(self, *args, **kwargs):
-        if not self.id and not self.site and not self.site_id:
+    def update_site_on_save(self, *args, **kwargs) -> None:
+        if not self.id:
             self.site = self.get_site_on_create()
         elif "update_fields" in kwargs and "site" not in kwargs.get("update_fields"):
             pass
         else:
             self.validate_site_against_current()
+        return None
 
     def get_site_on_create(self) -> Site:
         """Returns a site model instance.
 
         See also django-multisite.
         """
-        site = None
-        if not self.site:
-            try:
-                with transaction.atomic():
-                    site = get_site_model_cls().objects.get_current()
-            except ObjectDoesNotExist as e:
-                site_ids = [str(s) for s in sites.all()]
-                raise SiteModelMixinError(
-                    "Exception raised when trying manager method `get_current()`. "
-                    f"Sites registered with `sites` global are {site_ids}. "
-                    f"settings.SITE_ID={settings.SITE_ID}. Got {e}."
-                )
-        return site or self.site
+        try:
+            site_obj = self.site
+        except ObjectDoesNotExist:
+            site_obj = None
+        if not site_obj:
+            if self.site_id:
+                self.site = Site.objects.get(id=self.site_id)
+            else:
+                try:
+                    with transaction.atomic():
+                        site_obj = get_site_model_cls().objects.get_current()
+                except ObjectDoesNotExist as e:
+                    site_ids = [str(s) for s in sites.all()]
+                    raise SiteModelMixinError(
+                        "Exception raised when trying manager method `get_current()`. "
+                        f"Sites registered with `sites` global are {site_ids}. "
+                        f"settings.SITE_ID={settings.SITE_ID}. Got {e}."
+                    )
+        return site_obj
 
     def validate_site_against_current(self) -> None:
         """Validate existing site instance matches current_site."""
-        pass
+        return None
 
     class Meta:
         abstract = True
diff --git a/edc_sites/utils/valid_site_for_subject_or_raise.py b/edc_sites/utils/valid_site_for_subject_or_raise.py
index 1a962e2..c75ac95 100644
--- a/edc_sites/utils/valid_site_for_subject_or_raise.py
+++ b/edc_sites/utils/valid_site_for_subject_or_raise.py
@@ -30,18 +30,25 @@ def valid_site_for_subject_or_raise(
     )
     if skip_get_current_site:
         warn("Skipping validation of current site against registered subject site.")
-        current_site = registered_subject.site
+        site_obj = registered_subject.site
     else:
-        current_site: Site = get_site_model_cls().objects.get_current()
+        site_obj: Site = get_site_model_cls().objects.get_current()
         try:
-            registered_subject = get_registered_subject(
+            get_registered_subject(
                 subject_identifier,
-                site=current_site,
+                site=site_obj,
                 raise_exception=True,
             )
         except RegisteredSubjectDoesNotExist:
-            raise InvalidSiteForSubjectError(
-                f"Invalid site for subject. {subject_identifier}. "
-                f"Expected `{registered_subject.site.name}`. Got `{current_site.name}`"
-            )
-    return current_site
+            if not registered_subject.site_id:
+                raise InvalidSiteForSubjectError(
+                    "Site not defined for registered subject! "
+                    f"Subject identifier=`{subject_identifier}`. "
+                )
+            else:
+                raise InvalidSiteForSubjectError(
+                    f"Invalid site for subject. Subject identifier=`{subject_identifier}`. "
+                    f"Expected `{registered_subject.site.name}`. "
+                    f"Got site_id=`{site_obj.id}`"
+                )
+    return site_obj