From 062c8c7a5962ef23686ff1befa867287866c32d3 Mon Sep 17 00:00:00 2001
From: hanaawad24 <hana.awad@mcgovern.org>
Date: Wed, 27 Nov 2024 17:47:33 -0500
Subject: [PATCH] fix: failing test

---
 pvsite_datamodel/read/site.py |  2 +-
 pvsite_datamodel/sqlmodels.py | 14 ++------------
 tests/test_write.py           | 10 +++++-----
 3 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/pvsite_datamodel/read/site.py b/pvsite_datamodel/read/site.py
index c090f6e..68628d8 100644
--- a/pvsite_datamodel/read/site.py
+++ b/pvsite_datamodel/read/site.py
@@ -218,7 +218,7 @@ def assign_model_to_site(session: Session, model_uuid: str, site_uuid: str) -> s
     """
     site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()
 
-    model = session.query(MLModelSQL).filter(MLModelSQL.ml_model_id == model_uuid).first()
+    model = session.query(MLModelSQL).filter(MLModelSQL.model_uuid == model_uuid).first()
 
     if site is None:
         raise KeyError(f"Site with uuid {site_uuid} not found")
diff --git a/pvsite_datamodel/sqlmodels.py b/pvsite_datamodel/sqlmodels.py
index c7a039d..b8c3756 100644
--- a/pvsite_datamodel/sqlmodels.py
+++ b/pvsite_datamodel/sqlmodels.py
@@ -9,7 +9,6 @@
 from typing import List, Optional
 
 import sqlalchemy as sa
-from sqlalchemy import Column, ForeignKey, Integer
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import Mapped, declarative_base, relationship
 from sqlalchemy.schema import UniqueConstraint
@@ -37,8 +36,6 @@ class MLModelSQL(Base, CreatedMixin):
     )
     sites: Mapped[List["SiteSQL"]] = relationship("SiteSQL", back_populates="ml_model")
 
-    sites: Mapped[List["SiteSQL"]] = relationship("SiteSQL", back_populates="ml_model")
-
 
 class UserSQL(Base, CreatedMixin):
     """Class representing the users table.
@@ -185,13 +182,6 @@ class SiteSQL(Base, CreatedMixin):
         comment="The UUID of the client this site belongs to",
     )
 
-    ml_model_uuid = sa.Column(
-        UUID(as_uuid=True),
-        sa.ForeignKey("ml_model.model_uuid"),
-        nullable=True,
-        comment="The ML Model which should be used for this site",
-    )
-
     forecasts: Mapped[List["ForecastSQL"]] = relationship("ForecastSQL", back_populates="site")
     generation: Mapped[List["GenerationSQL"]] = relationship("GenerationSQL")
     inverters: Mapped[List["InverterSQL"]] = relationship(
@@ -203,8 +193,8 @@ class SiteSQL(Base, CreatedMixin):
     client: Mapped[List["ClientSQL"]] = relationship("ClientSQL", back_populates="sites")
     ml_model: Mapped[Optional[MLModelSQL]] = relationship("MLModelSQL", back_populates="sites")
 
-    ml_model_id = Column(Integer, ForeignKey("ml_models.id"))
-    ml_model = relationship("MLModel", back_populates="sites")
+    ml_model_uuid = sa.Column(UUID(as_uuid=True), sa.ForeignKey("ml_model.model_uuid"))
+    ml_model: Mapped[Optional[MLModelSQL]] = relationship("MLModelSQL", back_populates="sites")
 
 
 class ClientSQL(Base, CreatedMixin):
diff --git a/tests/test_write.py b/tests/test_write.py
index ecddc5b..ee67810 100644
--- a/tests/test_write.py
+++ b/tests/test_write.py
@@ -363,11 +363,11 @@ def test_assign_model_name_to_site(db_session):
 def test_assign_model_to_site(db_session):
     """Test to assign a model to a site"""
     site = make_fake_site(db_session=db_session)
-    model = get_or_create_model(session=db_session, model_name="Test Model")
+    model = get_or_create_model(session=db_session, name="Test Model")
 
-    message = assign_model_to_site(db_session, model.ml_model_uuid, site.site_uuid)
+    message = assign_model_to_site(db_session, model.model_uuid, site.site_uuid)
 
-    assert site.ml_model_uuid == model.ml_model_uuid
+    assert site.ml_model_uuid == model.model_uuid
     assert message == (
         f"Model {model.ml_model_uuid} successfully assigned to site {site.site_uuid}"
     )
@@ -375,11 +375,11 @@ def test_assign_model_to_site(db_session):
 
 def test_assign_model_to_nonexistent_site(db_session):
     """Test assigning a model to a nonexistent site"""
-    model = get_or_create_model(session=db_session, model_name="Test Model")
+    model = get_or_create_model(session=db_session, name="Test Model")
     nonexistent_site_uuid = str(uuid.uuid4())
 
     with pytest.raises(KeyError, match=f"Site with uuid {nonexistent_site_uuid} not found"):
-        assign_model_to_site(db_session, model.ml_model_uuid, nonexistent_site_uuid)
+        assign_model_to_site(db_session, model.model_uuid, nonexistent_site_uuid)
 
 
 def test_assign_nonexistent_model_to_site(db_session):