Add create_client and edit_client functions (openclimatefix#148)

added create_client and edit_client functions, and related tests.
edited PVSiteEditMetadata pydantic model, so edit_site function can update the client_uuid.
edited create_site function, so it can set the client_uuid.

README.md

@@ -76,6 +76,8 @@ Currently available write functions accessible via `from pvsite_datamodels.write
 - delete_user
 - delete_site_group
 - make_fake_site
+- create_client
+- edit_client
 
 
 ## Install the dependencies (requires [poetry][poetry])

pvsite_datamodel/pydantic_models.py

@@ -2,6 +2,7 @@
 
 from datetime import datetime
 from typing import Optional
+from uuid import UUID
 
 from pydantic import BaseModel, Field
 
@@ -63,3 +64,6 @@ class PVSiteEditMetadata(BaseModel):
     capacity_kw: Optional[float] = Field(None, description="The site's total capacity in kw", ge=0)
     dno: Optional[str] = Field(None, description="The site's DNO")
     gsp: Optional[str] = Field(None, description="The site's GSP")
+    client_uuid: Optional[UUID] = Field(
+        None, description="The UUID of the client this site belongs to"
+    )

pvsite_datamodel/write/__init__.py

@@ -7,12 +7,14 @@
 from .user_and_site import (
     add_site_to_site_group,
     change_user_site_group,
+    create_client,
     create_site,
     create_site_group,
     create_user,
     delete_site,
     delete_site_group,
     delete_user,
+    edit_client,
     edit_site,
     make_fake_site,
     update_user_site_group,

pvsite_datamodel/write/user_and_site.py

@@ -3,6 +3,7 @@
 import json
 from datetime import datetime, timezone
 from typing import Optional, Tuple
+from uuid import UUID
 
 import sqlalchemy as sa
 from sqlalchemy.orm.session import Session
@@ -11,6 +12,7 @@
 from pvsite_datamodel.pydantic_models import PVSiteEditMetadata
 from pvsite_datamodel.read import get_user_by_email
 from pvsite_datamodel.sqlmodels import (
+    ClientSQL,
     ForecastSQL,
     ForecastValueSQL,
     SiteAssetType,
@@ -74,6 +76,7 @@ def create_site(
     tilt: Optional[float] = None,
     inverter_capacity_kw: Optional[float] = None,
     module_capacity_kw: Optional[float] = None,
+    client_uuid: Optional[UUID] = None,
 ) -> [SiteSQL, str]:
     """
     Create a site and adds it to the database.
@@ -148,6 +151,7 @@ def create_site(
         tilt=tilt,
         inverter_capacity_kw=inverter_capacity_kw,
         module_capacity_kw=module_capacity_kw,
+        client_uuid=client_uuid,
     )
 
     session.add(site)
@@ -267,6 +271,7 @@ def edit_site(
         - capacity_kw: capacity of site in kw
         - dno: dno of site
         - gsp: gsp of site
+        - client_uuid: The UUID of the client this site belongs to
     """
     site = session.query(SiteSQL).filter(SiteSQL.site_uuid == site_uuid).first()
 
@@ -370,3 +375,35 @@ def delete_site_group(session: Session, site_group_name: str) -> str:
     session.commit()
 
     return message
+
+
+def create_client(session: Session, client_name: str) -> ClientSQL:
+    """Create a client.
+
+    :param session: database session
+    :param client_name: name of client being created
+    """
+    client = ClientSQL(client_name=client_name)
+
+    session.add(client)
+    session.commit()
+    session.refresh(client)
+
+    return client
+
+
+def edit_client(session: Session, client_uuid: UUID, client_name: str) -> ClientSQL:
+    """Edit an existing client.
+
+    :param session: database session
+    :param client_uuid: the existing client uuid
+    :param client_name: name of the client
+    """
+    client = session.query(ClientSQL).filter(ClientSQL.client_uuid == client_uuid).first()
+
+    client.client_name = client_name
+
+    session.add(client)
+    session.commit()
+
+    return client

tests/test_write.py

@@ -18,9 +18,11 @@
 from pvsite_datamodel.write.user_and_site import (
     add_site_to_site_group,
     change_user_site_group,
+    create_client,
     create_site,
     create_site_group,
     create_user,
+    edit_client,
     edit_site,
     make_fake_site,
 )
@@ -279,3 +281,42 @@ def test_edit_site(db_session):
     assert site.tilt == metadata_to_update.tilt
     assert site.capacity_kw == metadata_to_update.capacity_kw
     assert site.latitude == prev_latitude
+
+
+def test_create_client(db_session):
+    "Test to create a new client."
+    client = create_client(session=db_session, client_name="Test Client")
+
+    assert client.client_name == "Test Client"
+
+
+def test_edit_client(db_session):
+    "Test to edit a client."
+    client = create_client(session=db_session, client_name="Test Client")
+
+    client = edit_client(
+        session=db_session,
+        client_uuid=client.client_uuid,
+        client_name="Edited Client",
+    )
+
+    assert client.client_name == "Edited Client"
+
+
+def test_create_new_site_with_client_uuid(db_session):
+    "Test the creation of a site with a client uuid"
+    # creating a fake client owning the site
+    client = create_client(session=db_session, client_name="Test Client")
+
+    site, _ = create_site(
+        session=db_session,
+        client_site_id=6932,
+        client_site_name="test_site_name",
+        latitude=51.0,
+        longitude=0.0,
+        capacity_kw=1.0,
+        client_uuid=client.client_uuid,
+    )
+
+    assert site.client_site_name == "test_site_name"
+    assert site.client_uuid == client.client_uuid