added unit tests

localstack · Oct 11, 2024 · 29ba736 · 29ba736
1 parent 8aa20a2
commit 29ba736
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/moto/utilities/id_generator.py b/moto/utilities/id_generator.py
@@ -66,13 +66,13 @@ def __init__(self) -> None:
         self._id_sources = []
 
         self.add_id_source(self.get_id_from_tags)
-        self.add_id_source(self.get_custom_id)
+        self.add_id_source(self.get_custom_id_from_context)
 
-    def get_custom_id(self, id_source_context: IdSourceContext) -> Union[str, None]:
+    def get_custom_id(
+        self, resource_identifier: ResourceIdentifier
+    ) -> Union[str, None]:
         # retrieves a custom_id for a resource. Returns None
-        if resource_identifier := id_source_context.get("resource_identifier"):
-            return self._custom_ids.get(resource_identifier.unique_identifier)
-        return None
+        return self._custom_ids.get(resource_identifier.unique_identifier)
 
     def set_custom_id(
         self, resource_identifier: ResourceIdentifier, custom_id: str
@@ -100,6 +100,14 @@ def get_id_from_tags(id_source_context: IdSourceContext) -> Union[str, None]:
 
         return None
 
+    def get_custom_id_from_context(
+        self, id_source_context: IdSourceContext
+    ) -> Union[str, None]:
+        # retrieves a custom_id for a resource. Returns None
+        if resource_identifier := id_source_context.get("resource_identifier"):
+            return self.get_custom_id(resource_identifier)
+        return None
+
     def find_id_from_sources(
         self, id_source_context: IdSourceContext
     ) -> Union[str, None]:

diff --git a/tests/test_utilities/test_id_generator.py b/tests/test_utilities/test_id_generator.py
@@ -4,6 +4,7 @@
     ResourceIdentifier,
     Tags,
     moto_id,
+    moto_id_manager,
 )
 
 ACCOUNT = "account"
@@ -110,3 +111,24 @@ def test_generate_with_tags_fallback(set_custom_id):
         tags={TAG_KEY_CUSTOM_ID: TAG_ID},
     )
     assert generated_id == CUSTOM_ID
+
+
+def test_set_custom_id_lifecycle():
+    resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME)
+
+    moto_id_manager.set_custom_id(resource_identifier, CUSTOM_ID)
+
+    found_id = moto_id_manager.get_custom_id(resource_identifier)
+    assert found_id == CUSTOM_ID
+
+    moto_id_manager.unset_custom_id(resource_identifier)
+
+    found_id = moto_id_manager.get_custom_id(resource_identifier)
+    assert found_id is None
+
+
+def test_set_custom_id_name_is_not_set():
+    resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, None)
+    moto_id_manager.set_custom_id(resource_identifier, CUSTOM_ID)
+
+    assert moto_id_manager._custom_ids == {}