fixes async issues

RocketPy-Team · Apr 30, 2024 · 39dd2cd · 39dd2cd
1 parent d0d1fe3
commit 39dd2cd
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 43 deletions.
diff --git a/lib/controllers/environment.py b/lib/controllers/environment.py
@@ -69,9 +69,10 @@ async def create_env(self) -> "Union[EnvCreated, HTTPException]":
         Returns:
             views.EnvCreated
         """
-        env_repo = EnvRepository(environment=self.env)
         try:
-            await env_repo.create_env()
+            created_env = await EnvRepository(
+                environment=self.env
+            ).create_env()
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(
@@ -82,7 +83,7 @@ async def create_env(self) -> "Union[EnvCreated, HTTPException]":
                 detail=f"Failed to create environment: {e}",
             ) from e
         else:
-            return EnvCreated(env_id=env_repo.env_id)
+            return EnvCreated(env_id=created_env.env_id)
         finally:
             logging.info(
                 f"[{datetime.now()}] Call to controllers.environment.create_env completed; params: Env {hash(self.env)}"
@@ -102,9 +103,8 @@ async def get_env_by_id(env_id: int) -> "Union[Env, HTTPException]":
         Raises:
             HTTP 404 Not Found: If the env is not found in the database.
         """
-        env_repo = EnvRepository(env_id=env_id)
         try:
-            read_env = await env_repo.get_env()
+            read_env = await EnvRepository(env_id=env_id).get_env()
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(
@@ -143,9 +143,10 @@ async def get_rocketpy_env_as_jsonpickle(
         Raises:
             HTTP 404 Not Found: If the env is not found in the database.
         """
-        env_repo = EnvRepository(env_id=env_id)
         try:
-            read_env = await env_repo.get_env()
+            read_env = await cls.get_env_by_id(env_id)
+        except HTTPException as e:
+            raise e from e
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(
@@ -156,14 +157,9 @@ async def get_rocketpy_env_as_jsonpickle(
                 detail=f"Failed to read environment: {e}",
             ) from e
         else:
-            if read_env:
-                rocketpy_env = cls.get_rocketpy_env(read_env)
-                return EnvPickle(
-                    jsonpickle_rocketpy_env=jsonpickle.encode(rocketpy_env)
-                )
-            raise HTTPException(
-                status_code=status.HTTP_404_NOT_FOUND,
-                detail="Environment not found",
+            rocketpy_env = await cls.get_rocketpy_env(read_env)
+            return EnvPickle(
+                jsonpickle_rocketpy_env=jsonpickle.encode(rocketpy_env)
             )
         finally:
             logging.info(
@@ -185,16 +181,13 @@ async def update_env(
         Raises:
             HTTP 404 Not Found: If the env is not found in the database.
         """
-        env_repo = EnvRepository(environment=self.env, env_id=env_id)
         try:
-            read_env = await env_repo.get_env()
-            if read_env:
-                await env_repo.update_env()
-            else:
-                raise HTTPException(
-                    status_code=status.HTTP_404_NOT_FOUND,
-                    detail="Environment not found",
-                )
+            await EnvController.get_env_by_id(env_id)
+            updated_env = await EnvRepository(
+                environment=self.env, env_id=env_id
+            ).update_env()
+        except HTTPException as e:
+            raise e from e
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(
@@ -205,7 +198,7 @@ async def update_env(
                 detail=f"Failed to update environment: {e}",
             ) from e
         else:
-            return EnvUpdated(new_env_id=env_repo.env_id)
+            return EnvUpdated(new_env_id=updated_env.env_id)
         finally:
             logging.info(
                 f"[{datetime.now()}] Call to controllers.environment.update_env completed; params: EnvID {env_id}, Env {hash(self.env)}"
@@ -225,9 +218,8 @@ async def delete_env(env_id: str) -> "Union[EnvDeleted, HTTPException]":
         Raises:
             HTTP 404 Not Found: If the env is not found in the database.
         """
-        env_repo = EnvRepository(env_id=env_id)
         try:
-            await env_repo.delete_env()
+            await EnvRepository(env_id=env_id).delete_env()
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(
@@ -258,9 +250,9 @@ async def simulate(cls, env_id: int) -> "Union[EnvSummary, HTTPException]":
         Raises:
             HTTP 404 Not Found: If the env does not exist in the database.
         """
-        read_env = await cls.get_env_by_id(env_id)
         try:
-            rocketpy_env = cls.get_rocketpy_env(read_env)
+            read_env = await cls.get_env_by_id(env_id)
+            rocketpy_env = await cls.get_rocketpy_env(read_env)
             env_simulation_numbers = EnvData.parse_obj(
                 rocketpy_env.all_info_returned()
             )
@@ -270,6 +262,8 @@ async def simulate(cls, env_id: int) -> "Union[EnvSummary, HTTPException]":
             env_summary = EnvSummary(
                 env_data=env_simulation_numbers, env_plots=env_simulation_plots
             )
+        except HTTPException as e:
+            raise e from e
         except Exception as e:
             exc_str = parse_error(e)
             logging.error(

diff --git a/lib/repositories/environment.py b/lib/repositories/environment.py
@@ -22,7 +22,7 @@ def __init__(self, environment: Env = None, env_id: str = None):
         if env_id:
             self._env_id = env_id
         else:
-            self._env_id = hash(self._env)
+            self._env_id = str(hash(self._env))
 
     @property
     def env(self) -> "Env":
@@ -45,12 +45,8 @@ async def create_env(self):
         Creates a non-existing models.Env in the database
 
         Returns:
-            None
+            self
         """
-        env_exists = await self.get_env()
-        if env_exists:
-            return
-
         try:
             environment_to_dict = self.env.dict()
             environment_to_dict["env_id"] = self.env_id
@@ -61,22 +57,24 @@ async def create_env(self):
                 f"[{datetime.now()}] repositories.environment.create_env: {exc_str}"
             )
             raise Exception(f"Error creating environment: {str(e)}") from e
+        else:
+            return self
         finally:
             logging.info(
                 f"[{datetime.now()}] Call to repositories.environment.create_env completed; states: Env {hash(self.env)}, EnvID {self.env_id}"
             )
-            self.close_connection()
+            await self.close_connection()
 
     async def update_env(self):
         """
         Updates a models.Env in the database
 
         Returns:
-            None
+           self
         """
         try:
             environment_to_dict = self.env.dict()
-            environment_to_dict["env_id"] = hash(self.env)
+            environment_to_dict["env_id"] = str(hash(self.env))
             await self.collection.update_one(
                 {"env_id": self.env_id}, {"$set": environment_to_dict}
             )
@@ -88,12 +86,12 @@ async def update_env(self):
             )
             raise Exception(f"Error updating environment: {str(e)}") from e
         else:
-            return
+            return self
         finally:
             logging.info(
                 f"[{datetime.now()}] Call to repositories.environment.update_env completed; states: Env {hash(self.env)}, EnvID {self.env_id}"
             )
-            self.close_connection()
+            await self.close_connection()
 
     async def get_env(self) -> "Union[Env, None]":
         """
@@ -115,7 +113,7 @@ async def get_env(self) -> "Union[Env, None]":
             logging.info(
                 f"[{datetime.now()}] Call to repositories.environment.get_env completed; states: Env {hash(self.env)}, EnvID {self.env_id}"
             )
-            self.close_connection()
+            await self.close_connection()
 
     async def delete_env(self):
         """
@@ -135,4 +133,4 @@ async def delete_env(self):
             logging.info(
                 f"[{datetime.now()}] Call to repositories.environment.delete_env completed; states: Env {hash(self.env)}, EnvID {self.env_id}"
             )
-            self.close_connection()
+            await self.close_connection()
diff --git a/lib/repositories/repo.py b/lib/repositories/repo.py
@@ -27,6 +27,6 @@ def __init__(self, collection: str):
         self.db = self.client.rocketpy
         self.collection = self.db[collection]
 
-    def close_connection(self) -> None:
+    async def close_connection(self) -> None:
         logging.info(f"[{datetime.now()}] Closing connection to database")
         self.client.close()