Added user referenc to city creation

exports/db_factory.py

@@ -31,6 +31,13 @@ class DBFactory:
     """
     return self._city_repo.get_by_name(city_name)
 
+  def get_city_by_user(self, user_id):
+    """
+    Retrieve cities created by user
+    :param user_id: the id of the user
+    """
+    return self._city_repo.get_by_user(user_id)
+
   def get_hp_simulation(self, hp_sim_id: int):
     """
     Retrieve a single heat pump simulation from postgres

imports/db_factory.py

@@ -19,11 +19,11 @@ class DBFactory:
     self._city_repo = CityRepo(db_name=db_name, dotenv_path=dotenv_path, app_env=app_env)
     self._hp_simulation_repo = HeatPumpSimulationRepo(db_name=db_name, dotenv_path=dotenv_path, app_env=app_env)
 
-  def persist_city(self):
+  def persist_city(self, user_id: int):
     """
     Persist city into postgres database
     """
-    return self._city_repo.insert(self._city)
+    return self._city_repo.insert(self._city, user_id)
 
   def update_city(self, city_id, city):
     """

persistence/models/city.py

@@ -35,8 +35,9 @@ class City(Base):
   created = Column(DateTime, default=datetime.datetime.utcnow)
   updated = Column(DateTime, default=datetime.datetime.utcnow)
 
-  def __init__(self, city, name, srs_name, country_code, l_corner, u_corner):
+  def __init__(self, city, name, srs_name, country_code, l_corner, u_corner, user_id):
     self.city = city
+    self.user_id = user_id
     self.name = name
     self.srs_name = srs_name
     self.country_code = country_code

persistence/repositories/city_repo.py

@@ -32,9 +32,9 @@ class CityRepo(BaseRepo):
       cls._instance = super(CityRepo, cls).__new__(cls)
     return cls._instance
 
-  def insert(self, city: City) -> Union[City, Dict]:
+  def insert(self, city: City, user_id: int) -> Union[City, Dict]:
     db_city = DBCity(pickle.dumps(city), city.name, city.srs_name, city.country_code, city.lower_corner,
-                     city.upper_corner)
+                     city.upper_corner, user_id)
     db_city.climate_reference_city = city.climate_reference_city
     db_city.longitude = city.longitude
     db_city.latitude = city.latitude
@@ -127,6 +127,18 @@ class CityRepo(BaseRepo):
     except SQLAlchemyError as err:
       logger.error(f'Error while fetching city by name: {err}')
 
+  def get_by_user(self, user_id: int) -> [DBCity]:
+      """
+      Fetch city based on the user who created it
+      :param user_id: the id of the user
+      :return: [ModelCity] with the provided name
+      """
+      try:
+        result_set = self.session.execute(select(DBCity).where(DBCity.user_id == user_id))
+        return [building[0] for building in result_set]
+      except SQLAlchemyError as err:
+        logger.error(f'Error while fetching city by name: {err}')
+
   def delete_city(self, city_id: int):
     """
     Deletes a City with the id

unittests/test_db_factory.py

@@ -7,10 +7,12 @@ Project Coder Peter Yefi peteryefi@gmail.com
 from unittest import TestCase
 from imports.geometry_factory import GeometryFactory
 from imports.db_factory import DBFactory
+from imports.user_factory import UserFactory
 from exports.db_factory import DBFactory as ExportDBFactory
 from persistence.base_repo import BaseRepo
 from sqlalchemy import create_engine
 from persistence.models import City
+from persistence.models import User, UserRoles
 from pickle import loads
 from sqlalchemy.exc import ProgrammingError
 
@@ -43,16 +45,18 @@ class TestDBFactory(TestCase):
     cnn.execute('commit')
     cnn.execute("CREATE DATABASE test_db")
     cnn.close()
-
+    User.__table__.create(bind=repo.engine, checkfirst=True)
     City.__table__.create(bind=repo.engine, checkfirst=True)
 
     city_file = "../unittests/tests_data/C40_Final.gml"
     cls.city = GeometryFactory('citygml', city_file).city
     cls._db_factory = DBFactory(city=cls.city, db_name='test_db', app_env='TEST', dotenv_path='../.env')
     cls._export_db_factory = ExportDBFactory(db_name='test_db', app_env='TEST', dotenv_path='../.env')
+    user_factory = UserFactory(db_name='test_db', app_env='TEST', dotenv_path='../.env')
+    cls._user = user_factory.create_user("Admin", "admin@hub.com", "Admin@123", UserRoles.Admin)
 
   def test_save_city(self):
-    saved_city = self._db_factory.persist_city()
+    saved_city = self._db_factory.persist_city(self._user.id)
     self.assertEqual(saved_city.name, 'Montréal')
     pickled_city = loads(saved_city.city)
     self.assertEqual(len(pickled_city.buildings), 10)
@@ -60,27 +64,33 @@ class TestDBFactory(TestCase):
     self._db_factory.delete_city(saved_city.id)
 
   def test_save_same_city_with_same_hub_version(self):
-    first_city = self._db_factory.persist_city()
-    second_city = self._db_factory.persist_city()
+    first_city = self._db_factory.persist_city(self._user.id)
+    second_city = self._db_factory.persist_city(self._user.id)
     self.assertEqual(second_city['message'], f'Same version of {self.city.name} exist')
     self.assertEqual(first_city.name, 'Montréal')
     self.assertEqual(first_city.country_code, 'ca')
     self._db_factory.delete_city(first_city.id)
 
   def test_get_city_by_name(self):
-    city = self._db_factory.persist_city()
+    city = self._db_factory.persist_city(self._user.id)
     retrieved_city = self._export_db_factory.get_city_by_name(city.name)
     self.assertEqual(retrieved_city[0].lower_corner[0], 610610.7547462888)
     self._db_factory.delete_city(city.id)
 
+  def test_get_city_by_user(self):
+    city = self._db_factory.persist_city(self._user.id)
+    retrieved_city = self._export_db_factory.get_city_by_user(self._user.id)
+    self.assertEqual(retrieved_city[0].user_id, self._user.id)
+    self._db_factory.delete_city(city.id)
+
   def test_get_city_by_id(self):
-    city = self._db_factory.persist_city()
+    city = self._db_factory.persist_city(self._user.id)
     retrieved_city = self._export_db_factory.get_city(city.id)
     self.assertEqual(retrieved_city.upper_corner[0], 610818.6731258357)
     self._db_factory.delete_city(city.id)
 
   def test_get_update_city(self):
-    city = self._db_factory.persist_city()
+    city = self._db_factory.persist_city(self._user.id)
     self.city.longitude = 1.43589
     self.city.latitude = -9.38928339
     self._db_factory.update_city(city.id, self.city)