Bedrock: Add appropriate URL's to intercept

getmoto · Apr 27, 2024 · 56aacb6 · 56aacb6
1 parent 1e0c613
commit 56aacb6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/moto/bedrock/urls.py b/moto/bedrock/urls.py
@@ -1,5 +1,6 @@
 """bedrock base URL and path."""
 
+from ..bedrockagent.responses import AgentsforBedrockResponse
 from .responses import BedrockResponse
 
 url_bases = [
@@ -8,4 +9,21 @@
 
 url_paths = {
     "{0}/.*$": BedrockResponse.dispatch,
+    "{0}/agents/?$": AgentsforBedrockResponse.dispatch,
+    "{0}/agents/(?P<agent_name>[^/]+)/$": AgentsforBedrockResponse.dispatch,
+    "{0}/custom-models$": BedrockResponse.dispatch,
+    "{0}/custom-models/(?P<modelIdentifier>[^/]+)$": BedrockResponse.dispatch,
+    "{0}/custom-models/(?P<arn_prefix>[^/]+)/(?P<jobIdentifier>[^/]+)$": BedrockResponse.dispatch,
+    "{0}/knowledgebases$": AgentsforBedrockResponse.dispatch,
+    "{0}/knowledgebases/(?P<kb_name>[^/]+)$": AgentsforBedrockResponse.dispatch,
+    "{0}/knowledgebases/(?P<kb_name>[^/]+)/$": AgentsforBedrockResponse.dispatch,
+    "{0}/listTagsForResource$": BedrockResponse.dispatch,
+    "{0}/logging/modelinvocations$": BedrockResponse.dispatch,
+    "{0}/model-customization-jobs$": BedrockResponse.dispatch,
+    "{0}/model-customization-jobs/(?P<jobIdentifier>[^/]+)$": BedrockResponse.dispatch,
+    "{0}/model-customization-jobs/(?P<jobIdentifier>[^/]+)/stop$": BedrockResponse.dispatch,
+    "{0}/tags/(?P<resource_arn>[^/]+)$": AgentsforBedrockResponse.dispatch,
+    "{0}/tags/(?P<arn_prefix>[^/]+)/(?P<name>[^/]+)$": AgentsforBedrockResponse.dispatch,
+    "{0}/tagResource$": BedrockResponse.dispatch,
+    "{0}/untagResource$": BedrockResponse.dispatch,
 }
diff --git a/moto/bedrockagent/responses.py b/moto/bedrockagent/responses.py
@@ -132,21 +132,21 @@ def get_knowledge_base(self) -> str:
 
     def tag_resource(self) -> str:
         params = json.loads(self.body)
-        resource_arn = unquote(self.path.split("/")[-1])
+        resource_arn = unquote(self.path.split("/tags/")[-1])
         tags = params.get("tags")
         self.bedrockagent_backend.tag_resource(resource_arn=resource_arn, tags=tags)
         return json.dumps(dict())
 
     def untag_resource(self) -> str:
-        resource_arn = unquote(self.path.split("/")[-1])
+        resource_arn = unquote(self.path.split("/tags/")[-1])
         tag_keys = self.querystring.get("tagKeys", [])
         self.bedrockagent_backend.untag_resource(
             resource_arn=resource_arn, tag_keys=tag_keys
         )
         return json.dumps(dict())
 
     def list_tags_for_resource(self) -> str:
-        resource_arn = unquote(self.path.split("/")[-1])
+        resource_arn = unquote(self.path.split("/tags/")[-1])
         tags = self.bedrockagent_backend.list_tags_for_resource(
             resource_arn=resource_arn
         )