added tests

1 files changed, 14 insertions, 25 deletions

diff --git a/src/soc_collector/db.py b/src/soc_collector/db.py
index b10d865..b1501d8 100644
--- a/src/soc_collector/db.py
+++ b/src/soc_collector/db.py
@@ -7,7 +7,6 @@ from dataclasses import dataclass
 from fastapi import HTTPException
 from pydantic import BaseModel
 from bson import ObjectId
-from bson.errors import InvalidId
 from pymongo.errors import OperationFailure
 from pymongo import (
     ASCENDING,
@@ -19,19 +18,6 @@ from motor.motor_asyncio import (
 )
 
 
-def object_id_from_key(key: str) -> ObjectId:
-    """Get ObjectId from key, 400 if invalid ObjectId
-
-    :param key: Key.
-    :return: ObjectId
-    """
-
-    try:
-        return ObjectId(key)
-    except InvalidId as exc:
-        raise HTTPException(status_code=400, detail="Invalid key/object id") from exc
-
-
 class SearchInput(BaseModel):
     """Handle search data for HTTP request"""
 
@@ -71,7 +57,10 @@ class DBClient:
                 print(f"WARNING failed to connect to DB - {i} / 4", flush=True)
             sleep(1)
         else:
-            print("Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production")
+            print(
+                "Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production",
+                flush=True,
+            )
             app_exit(1)
 
     async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]:
@@ -82,11 +71,11 @@ class DBClient:
         """
 
         data: List[Dict[str, Any]] = []
-        cursor = self.collection.find(search_data.filter)
-
-        cursor.sort([("ip", ASCENDING), ("timestamp", DESCENDING)]).limit(search_data.limit).skip(search_data.skip)
 
         try:
+            cursor = self.collection.find(search_data.filter)
+            cursor.sort([("ip", ASCENDING), ("timestamp", DESCENDING)]).limit(search_data.limit).skip(search_data.skip)
+
             async for document in cursor:
                 if document is not None:
                     document["_id"] = str(document["_id"])
@@ -101,7 +90,7 @@ class DBClient:
                 detail="Probably wrong syntax, note the dictionary for find: "
                 + "https://motor.readthedocs.io/en/stable/tutorial-asyncio.html#async-for",
             ) from exc
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
@@ -119,7 +108,7 @@ class DBClient:
                 return document
             return None
 
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
@@ -136,7 +125,7 @@ class DBClient:
                 return result.inserted_id
             return None
 
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
@@ -154,7 +143,7 @@ class DBClient:
                 return object_id
             return None
 
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
@@ -171,7 +160,7 @@ class DBClient:
                 return object_id
             return None
 
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
@@ -182,11 +171,11 @@ class DBClient:
         """
 
         try:
-            result = await self.collection.estimated_document_count()
+            result = await self.collection.estimated_document_count(maxTimeMS=4000)
             if isinstance(result, int):
                 return result
             return None
 
-        except BaseException as exc:
+        except Exception as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc