fix: use names as ids for convenience

ZettaAI · Aug 22, 2024 · c1f9a89 · c1f9a89
1 parent 712d14d
commit c1f9a89
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 11 deletions.
diff --git a/tests/unit/db_annotations/test_annotation.py b/tests/unit/db_annotations/test_annotation.py
@@ -7,16 +7,17 @@
 from zetta_utils.db_annotations import annotation, collection, layer, layer_group
 
 
-def _init_collection_and_layer_group() -> tuple[str, str]:
+def _init_collection_and_layer_group(
+    collection_name: str, layer_group_name: str
+) -> tuple[str, str]:
     user = "john_doe"
-    collection_name = "test_ann_collection0"
     collection_id = collection.add_collection(collection_name, user, "this is a test")
 
     layer_id0 = layer.add_layer("test_layer0", "precomputed://test0", "this is a test")
     layer_id1 = layer.add_layer("test_layer1", "precomputed://test1", "this is a test")
 
     layer_group_id = layer_group.add_layer_group(
-        name="test_layer_group0",
+        name=layer_group_name,
         collection_id=collection_id,
         user=user,
         layers=[layer_id0, layer_id1],
@@ -28,7 +29,9 @@ def _init_collection_and_layer_group() -> tuple[str, str]:
 def test_add_update_delete_annotation(
     firestore_emulator, annotations_db, collections_db, layer_groups_db, layers_db
 ):
-    collection_id, layer_group_id = _init_collection_and_layer_group()
+    collection_id, layer_group_id = _init_collection_and_layer_group(
+        "new_collection0", "new_lgroup0"
+    )
     annotation_raw = {
         "pointA": [1, 1, 1],
         "pointB": [1, 1, 5],
@@ -67,7 +70,9 @@ def test_add_update_delete_annotation(
 def test_add_update_annotations(
     firestore_emulator, annotations_db, collections_db, layer_groups_db, layers_db
 ):
-    collection_id, layer_group_id = _init_collection_and_layer_group()
+    collection_id, layer_group_id = _init_collection_and_layer_group(
+        "new_collection1", "new_lgroup1"
+    )
     annotations_raw = [
         {
             "pointA": [1, 1, 1],
@@ -120,7 +125,9 @@ def test_add_update_annotations(
 
 
 def test_read_delete_annotations(firestore_emulator, annotations_db):
-    collection_id, layer_group_id = _init_collection_and_layer_group()
+    collection_id, layer_group_id = _init_collection_and_layer_group(
+        "new_collection2", "new_lgroup2"
+    )
     ng_annotations = annotation.parse_ng_annotations(
         [
             {

diff --git a/tests/unit/db_annotations/test_layer_group.py b/tests/unit/db_annotations/test_layer_group.py
@@ -46,7 +46,7 @@ def test_add_update_delete_layer_group(firestore_emulator, layer_groups_db):
 
 def test_read_delete_layer_groups(firestore_emulator, layer_groups_db):
     user = "john_doe"
-    collection_name = "test_lg_collection0"
+    collection_name = "test_lg_collection1"
     collection_id = collection.add_collection(collection_name, user, "this is a test")
 
     layer_id0 = layer.add_layer("test_layer0", "precomputed://test0", "this is a test")

diff --git a/zetta_utils/db_annotations/collection.py b/zetta_utils/db_annotations/collection.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import time
-import uuid
 from typing import overload
 
 from zetta_utils.layer.db_layer import DBRowDataT
@@ -47,7 +46,9 @@ def read_collections(
 
 
 def add_collection(name: str, user: str, comment: str | None = None) -> str:
-    collection_id = str(uuid.uuid4())
+    collection_id = name
+    if collection_id in COLLECTIONS_DB:
+        raise KeyError(f"{collection_id} already exists.")
     col_keys = INDEXED_COLS + NON_INDEXED_COLS
     row: DBRowDataT = {"name": name, "created_by": user, "created_at": time.time()}
     if comment:

diff --git a/zetta_utils/db_annotations/layer_group.py b/zetta_utils/db_annotations/layer_group.py
@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import uuid
 from typing import overload
 
 from zetta_utils.layer.db_layer import DBRowDataT
@@ -62,7 +61,9 @@ def add_layer_group(
     layers: list[str] | None = None,
     comment: str | None = None,
 ) -> str:
-    layer_group_id = str(uuid.uuid4())
+    layer_group_id = name
+    if layer_group_id in LAYER_GROUPS_DB:
+        raise KeyError(f"{layer_group_id} already exists.")
     col_keys = INDEXED_COLS + NON_INDEXED_COLS
     row: DBRowDataT = {"name": name, "collection": collection_id, "created_by": user}
     if layers: