1 files changed, 9 insertions, 9 deletions
diff --git a/src/collector/db.py b/src/collector/db.py
index 54cdc21..641e5da 100644
--- a/src/collector/db.py
+++ b/src/collector/db.py
@@ -56,7 +56,7 @@ class DBClient:
             print("Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production")
             app_exit(1)
 
-    async def find(self, search_data: SearchInput) -> Optional[List[Dict[str, Any]]]:
+    async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]:
         """Wrap the find() method, handling timeouts and return data type.
 
         :param search_data: Instance of SearchInput.
@@ -67,14 +67,16 @@ class DBClient:
         cursor = self.collection.find(search_data.search)
 
         # Sort on timestamp
+        # TODO: Also sort on IP as well
         cursor.sort("timestamp", -1).limit(search_data.limit).skip(search_data.skip)
 
         try:
             async for document in cursor:
-                data.append(document)
+                if document is not None:
+                    document["_id"] = str(document["_id"])
+                    data.append(document)
 
-            if data:
-                return data
+            return data
 
         except OperationFailure as exc:
             print(f"DB failed to process: {exc.details}")
@@ -87,8 +89,6 @@ class DBClient:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
-        return None
-
     async def find_one(self, object_id: ObjectId) -> Optional[Dict[str, Any]]:
         """Wrap the find_one() method, handling timeouts and return data type.
 
@@ -99,6 +99,7 @@ class DBClient:
         try:
             document = await self.collection.find_one({"_id": object_id})
             if isinstance(document, Dict):
+                document["_id"] = str(document["_id"])
                 return document
             return None
 
@@ -150,15 +151,14 @@ class DBClient:
 
         try:
             result = await self.collection.delete_one({"_id": object_id})
-            if result.deleted_count == 1:
+            if isinstance(result.deleted_count, int) and result.deleted_count == 1:
                 return object_id
+            return None
 
         except BaseException as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
-        return None
-
     async def estimated_document_count(self) -> Optional[int]:
         """Wrap the estimated_document_count() method, handling timeouts and return data type.