added tests

1 files changed, 27 insertions, 1 deletions

diff --git a/src/soc_collector/schema.py b/src/soc_collector/schema.py
index 221990a..2c2dfb9 100644
--- a/src/soc_collector/schema.py
+++ b/src/soc_collector/schema.py
@@ -1,6 +1,8 @@
 """Our schema module"""
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Union
 import jsonschema
+from bson import ObjectId
+from bson.errors import InvalidId
 
 # fmt:off
 # NOTE: Commented out properties are left intentionally, so it is easier to see
@@ -99,9 +101,33 @@ def valid_schema(json_data: Dict[str, Any]) -> bool:
     :param json_data: Json object
     :return: bool
     """
+
     try:
         jsonschema.validate(json_data, schema, format_checker=jsonschema.FormatChecker())
     except jsonschema.exceptions.ValidationError as exc:
         print(f"Validation failed with error: {exc.message}")
         return False
     return True
+
+
+def object_id_from_data(data: Union[str, Dict[str, Any]]) -> Optional[ObjectId]:
+    """Get ObjectId from key. None if invalid.
+
+    :param data: Key.
+    :return: Optional[ObjectId]
+    """
+
+    if isinstance(data, str):
+        try:
+            return ObjectId(data)
+        except InvalidId:
+            return None
+
+    elif isinstance(data, Dict):
+        if "_id" in data and isinstance(data["_id"], str):
+            try:
+                return ObjectId(data["_id"])
+            except InvalidId:
+                return None
+
+    return None