Add execute tool + scratch pad tests #4400
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
nathaliellenaa
requested review from
HenryL27,
Zhangxunmt,
austintlee,
b4sjoo,
dhrubo-os,
jngz-es,
mingshl,
model-collapse,
pyek-bot,
rbhavna,
sam-herman,
xinyual,
ylwu-amzn and
zane-neo
as code owners
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
mingshl
previously approved these changes
Nov 5, 2025
mingshl
reviewed
Nov 5, 2025
| super.setUp(); | ||
| } | ||
|
|
||
| public void testScratchpadSizeLimit() throws Exception { |
Collaborator
There was a problem hiding this comment.
thanks for verifying the exception in UT.
nathaliellenaa
force-pushed
the
tool_tests
branch
from
01c9af4 to
ddcaf7c
Compare
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
mingshl
reviewed
Nov 6, 2025
Collaborator
Seems like we started having a dependency conflict issue. |
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
ylwu-amzn
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
Contributor
Author
|
Failing tests are unrelated |
Collaborator
|
@nathaliellenaa this test is constantly failing, can you try fix the test by downloading another imageurl? we need the required CI passed to merge this. |
nathaliellenaa
force-pushed
the
tool_tests
branch
from
fcf1c94 to
e1f019e
Compare
nathaliellenaa
temporarily deployed
to
ml-commons-cicd-env-require-approval
GitHub Actions
Inactive
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
nathaliellenaa
temporarily deployed
to
ml-commons-cicd-env-require-approval
GitHub Actions
Inactive
5 tasks
jiapingzeng
approved these changes
Nov 26, 2025
rithin-pullela-aws
approved these changes
Nov 28, 2025
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
force-pushed
the
tool_tests
branch
from
e1f019e to
6ed156d
Compare
WalkthroughAdds unit and integration tests covering SecurityException/permission-denied paths and a scratchpad size-limit rejection; no production code changes. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Failure
nathaliellenaa
had a problem deploying
to
ml-commons-cicd-env-require-approval
GitHub Actions
Error
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/WriteToScratchPadToolTests.java (2)
215-225: Consider asserting scratchpad mutation in this JSON-array conversion test
testRun_StringConversion_WithJsonArraynow only checks the response string; asserting the updatedSCRATCHPAD_NOTES_KEYhere as well would better match the test name and keep this path fully covered (even though other tests cover similar mutations).
227-242: SecurityException test does not exerciseWriteToScratchPadToolbehavior
testRun_SecurityExceptiondirectly callslistener.onFailure(securityException)and then verifies the listener, without invokingtool.run(...). That means the test is not validating howWriteToScratchPadToolitself behaves when aSecurityExceptionoccurs; it only checks that Mockito and the listener work as expected.If the goal is tool-level behavior, consider:
- Driving
tool.run(...)and arranging dependencies so aSecurityExceptionflows through the listener, or- Dropping this test and relying on higher‑level tests like
MLToolExecutorTest.test_ToolExecutionFailsWithoutProperPermission, which already validate the permission error path end‑to‑end.ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java (1)
220-233: SecurityException test does not exerciseReadFromScratchPadToolbehavior
testRun_SecurityExceptionmanually triggerslistener.onFailure(securityException)and verifies the captured exception, but never callstool.run(...). This doesn’t assert anything about howReadFromScratchPadToolhandles permission failures; it just checks the listener wiring.To make this more meaningful, consider:
- Refactoring to drive
tool.run(...)and simulate aSecurityExceptionfrom its internal operations, or- Removing this test and relying on higher‑level permission tests (e.g.,
MLToolExecutorTestand the ITs) for SecurityException coverage.plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java (1)
44-80: Robust permission IT; only minor polish possibleThis test cleanly verifies that a no‑permission user gets a permission/forbidden/unauthorized error over HTTPS, and it correctly cleans up the client and user. The string‑contains checks are flexible enough to tolerate backend wording differences.
If you want to tighten things later, optional tweaks would be:
- Use
Collections.emptyList()instead ofnew ArrayList<>()for the no‑roles argument.- Use try‑with‑resources for
RestClientinstead of an explicittry/finally.plugin/src/test/java/org/opensearch/ml/tools/ScratchPadToolIT.java (1)
1-39: Effective size‑limit IT; consider tuning payload size / constructionThis IT correctly exercises the WriteToScratchPad tool via the REST endpoint with an oversized payload and asserts that some form of content‑length / 413‑style error is returned, without over‑fitting to a specific exception type or message.
Two optional refinements you might consider:
- Use a payload just over the configured
http.max_content_length(rather than a full 100 MB) to reduce memory pressure in CI while still triggering the same failure.- If you ever change
largeContentto something more complex than"A", switch to building the JSON body via a helper or escaping utility instead of rawString.formatto avoid quoting/escaping pitfalls.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutorTest.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/WriteToScratchPadToolTests.java(1 hunks)plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java(2 hunks)plugin/src/test/java/org/opensearch/ml/tools/ScratchPadToolIT.java(1 hunks)
🔇 Additional comments (1)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutorTest.java (1)
206-226: Good coverage of permission‑denied tool execution path
test_ToolExecutionFailsWithoutProperPermissioncorrectly simulates aSecurityExceptionfrom the tool and verifies it is propagated unchanged toactionListener, including the message fragment. This nicely complementstest_ToolExecutionFailedand the scratchpad tests.
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
nathaliellenaa
force-pushed
the
tool_tests
branch
from
6ed156d to
3bdb967
Compare
nathaliellenaa
requested a deployment
to
ml-commons-cicd-env-require-approval
GitHub Actions
Waiting
nathaliellenaa
requested a deployment
to
ml-commons-cicd-env-require-approval
GitHub Actions
Waiting
nathaliellenaa
requested a deployment
to
ml-commons-cicd-env-require-approval
GitHub Actions
Waiting
nathaliellenaa
requested a deployment
to
ml-commons-cicd-env-require-approval
GitHub Actions
Waiting
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java (1)
56-61: Consider reducing socket timeout for faster test execution.Permission checks typically fail quickly. The 60-second socket timeout could unnecessarily delay test failures if there are connectivity issues.
final RestClient noPermissionClient = new SecureRestClientBuilder( getClusterHosts().toArray(new HttpHost[0]), isHttps(), noPermissionUser, password - ).setSocketTimeout(60000).build(); + ).setSocketTimeout(10000).build();
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutorTest.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java(1 hunks)ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/WriteToScratchPadToolTests.java(1 hunks)plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java(2 hunks)plugin/src/test/java/org/opensearch/ml/tools/ScratchPadToolIT.java(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/WriteToScratchPadToolTests.java
- plugin/src/test/java/org/opensearch/ml/tools/ScratchPadToolIT.java
- ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutorTest.java
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: spotless
🔇 Additional comments (1)
plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java (1)
17-17: LGTM!The added imports are appropriate for the new permission test and are all utilized correctly.
Also applies to: 20-23
Comment on lines
+220
to
+233
| @Test | ||
| public void testRun_SecurityException() { | ||
| Map<String, String> parameters = new HashMap<>(); | ||
| parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"confidential data\"]"); | ||
|
|
||
| SecurityException securityException = new SecurityException("no permissions for [indices:data/read/get] and User [name=test_user]"); | ||
| listener.onFailure(securityException); | ||
|
|
||
| ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class); | ||
| verify(listener).onFailure(exceptionCaptor.capture()); | ||
| Exception exception = exceptionCaptor.getValue(); | ||
| assertTrue(exception instanceof SecurityException); | ||
| assertTrue(exception.getMessage().contains("no permissions")); | ||
| } |
There was a problem hiding this comment.
Test doesn't invoke production code.
This test directly calls listener.onFailure(securityException) without ever invoking tool.run(parameters, listener). It only verifies that Mockito captured the manually triggered failure, providing zero coverage of how ReadFromScratchPadTool actually handles a SecurityException.
To properly test SecurityException handling, you need to either:
- Mock an internal dependency that
tool.run()calls to throwSecurityException, or - Move this to an integration test where permissions are enforced
Apply this diff if you want to keep it as a unit test with mocked internals (note: this assumes the tool calls some external service that could throw SecurityException - adjust based on actual implementation):
@Test
public void testRun_SecurityException() {
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"confidential data\"]");
-
- SecurityException securityException = new SecurityException("no permissions for [indices:data/read/get] and User [name=test_user]");
- listener.onFailure(securityException);
-
+
+ // You need to mock whatever dependency throws SecurityException during tool.run()
+ // For example, if the tool reads from a secure store:
+ // when(mockSecureStore.read(any())).thenThrow(new SecurityException("no permissions..."));
+
+ tool.run(parameters, listener);
+
ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
verify(listener).onFailure(exceptionCaptor.capture());
Exception exception = exceptionCaptor.getValue();
assertTrue(exception instanceof SecurityException);
assertTrue(exception.getMessage().contains("no permissions"));
}Alternatively, if ReadFromScratchPadTool.run() doesn't actually interact with permission-protected resources, this test might be testing the wrong layer and should be removed or moved to an integration test.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In
ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java
around lines 220 to 233, the test currently calls listener.onFailure(...)
directly instead of exercising the production path; update the test to invoke
tool.run(parameters, listener) and cause a SecurityException from the tool's
dependency so the tool's failure handling is executed: identify the internal
dependency the tool calls (e.g., a client/service used to read the scratchpad),
mock that dependency to throw new SecurityException("no permissions...") when
invoked, inject the mock into the tool under test, call tool.run(parameters,
listener), then verify listener.onFailure(...) was called with a
SecurityException and assert its message; if the tool has no mockable
dependency, remove or convert this unit test into an integration test that
exercises real permission checks instead.
Comment on lines
+44
to
+80
| public void testListIndexWithNoPermissions() throws Exception { | ||
| if (!isHttps()) { | ||
| log.info("Skipping permission test as security is not enabled"); | ||
| return; | ||
| } | ||
|
|
||
| String noPermissionUser = "no_permission_user"; | ||
| String password = "TestPassword123!"; | ||
|
|
||
| try { | ||
| createUser(noPermissionUser, password, new ArrayList<>()); | ||
|
|
||
| final RestClient noPermissionClient = new SecureRestClientBuilder( | ||
| getClusterHosts().toArray(new HttpHost[0]), | ||
| isHttps(), | ||
| noPermissionUser, | ||
| password | ||
| ).setSocketTimeout(60000).build(); | ||
|
|
||
| try { | ||
| ResponseException exception = expectThrows(ResponseException.class, () -> { | ||
| TestHelper | ||
| .makeRequest(noPermissionClient, "POST", "/_plugins/_ml/agents/" + agentId + "/_execute", null, question, null); | ||
| }); | ||
|
|
||
| String errorMessage = exception.getMessage().toLowerCase(); | ||
| assertTrue( | ||
| "Expected permission error, got: " + errorMessage, | ||
| errorMessage.contains("no permissions") || errorMessage.contains("forbidden") || errorMessage.contains("unauthorized") | ||
| ); | ||
| } finally { | ||
| noPermissionClient.close(); | ||
| } | ||
| } finally { | ||
| deleteUser(noPermissionUser); | ||
| } | ||
| } |
There was a problem hiding this comment.
Missing @test annotation - test won't execute.
The test method testListIndexWithNoPermissions lacks the @Test annotation, so it will be silently skipped by the test runner.
Apply this diff:
+@Test
public void testListIndexWithNoPermissions() throws Exception {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| public void testListIndexWithNoPermissions() throws Exception { | |
| if (!isHttps()) { | |
| log.info("Skipping permission test as security is not enabled"); | |
| return; | |
| } | |
| String noPermissionUser = "no_permission_user"; | |
| String password = "TestPassword123!"; | |
| try { | |
| createUser(noPermissionUser, password, new ArrayList<>()); | |
| final RestClient noPermissionClient = new SecureRestClientBuilder( | |
| getClusterHosts().toArray(new HttpHost[0]), | |
| isHttps(), | |
| noPermissionUser, | |
| password | |
| ).setSocketTimeout(60000).build(); | |
| try { | |
| ResponseException exception = expectThrows(ResponseException.class, () -> { | |
| TestHelper | |
| .makeRequest(noPermissionClient, "POST", "/_plugins/_ml/agents/" + agentId + "/_execute", null, question, null); | |
| }); | |
| String errorMessage = exception.getMessage().toLowerCase(); | |
| assertTrue( | |
| "Expected permission error, got: " + errorMessage, | |
| errorMessage.contains("no permissions") || errorMessage.contains("forbidden") || errorMessage.contains("unauthorized") | |
| ); | |
| } finally { | |
| noPermissionClient.close(); | |
| } | |
| } finally { | |
| deleteUser(noPermissionUser); | |
| } | |
| } | |
| @Test | |
| public void testListIndexWithNoPermissions() throws Exception { | |
| if (!isHttps()) { | |
| log.info("Skipping permission test as security is not enabled"); | |
| return; | |
| } | |
| String noPermissionUser = "no_permission_user"; | |
| String password = "TestPassword123!"; | |
| try { | |
| createUser(noPermissionUser, password, new ArrayList<>()); | |
| final RestClient noPermissionClient = new SecureRestClientBuilder( | |
| getClusterHosts().toArray(new HttpHost[0]), | |
| isHttps(), | |
| noPermissionUser, | |
| password | |
| ).setSocketTimeout(60000).build(); | |
| try { | |
| ResponseException exception = expectThrows(ResponseException.class, () -> { | |
| TestHelper | |
| .makeRequest(noPermissionClient, "POST", "/_plugins/_ml/agents/" + agentId + "/_execute", null, question, null); | |
| }); | |
| String errorMessage = exception.getMessage().toLowerCase(); | |
| assertTrue( | |
| "Expected permission error, got: " + errorMessage, | |
| errorMessage.contains("no permissions") || errorMessage.contains("forbidden") || errorMessage.contains("unauthorized") | |
| ); | |
| } finally { | |
| noPermissionClient.close(); | |
| } | |
| } finally { | |
| deleteUser(noPermissionUser); | |
| } | |
| } |
🤖 Prompt for AI Agents
plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java around lines
44 to 80: the test method testListIndexWithNoPermissions is missing the @Test
annotation so it won’t be executed; add the appropriate @Test annotation above
the method and ensure the corresponding import (e.g., org.junit.Test or
org.junit.jupiter.api.Test matching the test framework used in this module) is
present at the top of the file.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Add execute tool + scratch pad tests
Related Issues
Resolves #[Issue number to be closed when this PR is merged]
Check List
--signoff.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.