Refactor assign_model_name_to_site function and add error handling fo… (

#186)

* Refactor assign_model_name_to_site function and add error handling for nonexistent site and test cases

* Fix linter

* Fix linter

* Remove not found test case in method and readapt exception message

pvsite_datamodel/write/user_and_site.py

@@ -379,8 +379,11 @@ def delete_site_group(session: Session, site_group_name: str) -> str:
 
 
 def assign_model_name_to_site(session: Session, site_uuid, model_name):
-    """
-    Assign a model name to a site.
+    """Assign model to site.
+
+    :param session: database session
+    :param site_uuid: site uuid
+    :param model_name: name of the model
     """
 
     site = get_site_by_uuid(session=session, site_uuid=site_uuid)

tests/write/test_write.py

@@ -356,3 +356,13 @@ def test_assign_model_name_to_site(db_session):
     assign_model_name_to_site(db_session, site.site_uuid, "test_model_2")
 
     assert site.ml_model.name == "test_model_2"
+
+
+def test_assign_model_to_nonexistent_site(db_session):
+    """Test assigning a model to a nonexistent site"""
+    nonexistent_site_uuid = str(uuid.uuid4())
+
+    with pytest.raises(
+        KeyError, match=f"Site uuid {nonexistent_site_uuid} not found in sites table"
+    ):
+        assign_model_name_to_site(db_session, nonexistent_site_uuid, "test_model")