From 3550786cb432e9bfdb454d7e54ef114f76782c5b Mon Sep 17 00:00:00 2001 From: Daniel Widdis Date: Thu, 23 Jan 2025 23:20:39 -0800 Subject: [PATCH] More code review comments and refactoring Signed-off-by: Daniel Widdis --- build.gradle | 6 +- .../indices/FlowFrameworkIndicesHandler.java | 11 +- .../rest/AbstractSearchWorkflowAction.java | 4 +- .../rest/RestCreateWorkflowAction.java | 4 +- .../rest/RestDeleteWorkflowAction.java | 4 +- .../rest/RestDeprovisionWorkflowAction.java | 4 +- .../rest/RestGetWorkflowAction.java | 4 +- .../rest/RestGetWorkflowStateAction.java | 4 +- .../rest/RestGetWorkflowStepAction.java | 4 +- .../rest/RestProvisionWorkflowAction.java | 4 +- .../flowframework/util/EncryptorUtils.java | 2 +- .../flowframework/util/RestActionUtils.java | 54 --------- .../flowframework/util/TenantAwareHelper.java | 34 +++++- .../FlowFrameworkTenantAwareRestTestCase.java | 4 +- .../rest/FlowFrameworkRestApiIT.java | 4 +- .../util/RestActionUtilsTests.java | 103 ----------------- .../util/TenantAwareHelperTests.java | 105 ++++++++++++++++-- 17 files changed, 156 insertions(+), 199 deletions(-) delete mode 100644 src/main/java/org/opensearch/flowframework/util/RestActionUtils.java delete mode 100644 src/test/java/org/opensearch/flowframework/util/RestActionUtilsTests.java diff --git a/build.gradle b/build.gradle index 1fcd8227d..f439c412b 100644 --- a/build.gradle +++ b/build.gradle @@ -192,9 +192,6 @@ dependencies { // Multi-tenant SDK Client implementation ("org.opensearch:opensearch-remote-metadata-sdk:${opensearch_build}") { exclude group: "org.apache.httpcomponents.client5", module: "httpclient5" - // Temporary pending merge of https://github.com/opensearch-project/opensearch-remote-metadata-sdk/pull/56 - exclude group: "org.apache.logging.log4j", module: "log4j-api" - exclude group: "org.apache.logging.log4j", module: "log4j-core" } testImplementation 'org.junit.jupiter:junit-jupiter:5.11.4' testImplementation "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}" @@ -281,7 +278,7 @@ integTest { if (System.getProperty("tests.rest.cluster") != null && System.getProperty("tests.rest.tenantaware") == null) { filter { includeTestsMatching "org.opensearch.flowframework.rest.*IT" - excludeTestsMatching "org.opensearch.ml.rest.*TenantAwareIT" + excludeTestsMatching "org.opensearch.flowframework.rest.*TenantAwareIT" } } @@ -304,6 +301,7 @@ integTest { filter { includeTestsMatching "org.opensearch.flowframework.rest.FlowFrameworkSecureRestApiIT" excludeTestsMatching "org.opensearch.flowframework.rest.FlowFrameworkRestApiIT" + excludeTestsMatching "org.opensearch.flowframework.rest.*TenantAwareIT" } } diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java index 7ef77966d..31c49bb14 100644 --- a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java +++ b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java @@ -475,7 +475,7 @@ public void updateTemplateInGlobalContext( documentId ).getFormattedMessage(); logger.error(errorMessage); - listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST)); + listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.INTERNAL_SERVER_ERROR)); return; } doesTemplateExist(documentId, tenantId, templateExists -> { @@ -792,15 +792,14 @@ public void updateFlowFrameworkSystemIndexDoc( sdkClient.updateDataObjectAsync(updateRequest).whenComplete((r, throwable) -> { context.restore(); if (throwable == null) { - UpdateResponse response; try { - response = UpdateResponse.fromXContent(r.parser()); - logger.info("Deleted workflow state doc: {}", documentId); + UpdateResponse response = UpdateResponse.fromXContent(r.parser()); + logger.info("Updated workflow state doc: {}", documentId); listener.onResponse(response); } catch (Exception e) { - logger.error("Failed to parse delete response", e); + logger.error("Failed to parse update response", e); listener.onFailure( - new FlowFrameworkException("Failed to parse delete response", RestStatus.INTERNAL_SERVER_ERROR) + new FlowFrameworkException("Failed to parse update response", RestStatus.INTERNAL_SERVER_ERROR) ); } } else { diff --git a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java index 52c0f7754..e5ae25287 100644 --- a/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java @@ -17,7 +17,7 @@ import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.flowframework.common.FlowFrameworkSettings; import org.opensearch.flowframework.exception.FlowFrameworkException; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; @@ -84,7 +84,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli new BytesRestResponse(ffe.getRestStatus(), ffe.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS)) ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchSourceBuilder.parseXContent(request.contentOrSourceParamParser()); searchSourceBuilder.seqNoAndPrimaryTerm(true).version(true); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java index e79e96319..a3beab336 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java @@ -25,7 +25,7 @@ import org.opensearch.flowframework.transport.CreateWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; import org.opensearch.flowframework.util.ParseUtils; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -115,7 +115,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli new BytesRestResponse(ffe.getRestStatus(), ffe.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS)) ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); if (!provision && !params.isEmpty()) { FlowFrameworkException ffe = new FlowFrameworkException( "Only the parameters " + request.consumedParams() + " are permitted unless the provision parameter is set to true.", diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java index 0ea4de8e0..65251d235 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestDeleteWorkflowAction.java @@ -20,7 +20,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.DeleteWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -73,7 +73,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); // Always consume content to silently ignore it // https://github.com/opensearch-project/flow-framework/issues/578 diff --git a/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java index 2f89fb8e2..783404ff0 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestDeprovisionWorkflowAction.java @@ -20,7 +20,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.DeprovisionWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -70,7 +70,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); // Always consume content to silently ignore it // https://github.com/opensearch-project/flow-framework/issues/578 diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java index ec08f4a09..3a05b4fb4 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowAction.java @@ -20,7 +20,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.GetWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -71,7 +71,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); // Always consume content to silently ignore it // https://github.com/opensearch-project/flow-framework/issues/578 request.content(); diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java index 7af1ce314..286d4d1e1 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStateAction.java @@ -20,7 +20,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.GetWorkflowStateAction; import org.opensearch.flowframework.transport.GetWorkflowStateRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -66,7 +66,7 @@ protected BaseRestHandler.RestChannelConsumer prepareRequest(RestRequest request RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); // Always consume content to silently ignore it // https://github.com/opensearch-project/flow-framework/issues/578 diff --git a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java index 1c5751653..d970a080d 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestGetWorkflowStepAction.java @@ -20,7 +20,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.GetWorkflowStepAction; import org.opensearch.flowframework.transport.WorkflowRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -71,7 +71,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkSettings.isMultiTenancyEnabled(), request); // Always consume content to silently ignore it // https://github.com/opensearch-project/flow-framework/issues/578 diff --git a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java index f3d2977ca..fa53aec27 100644 --- a/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java +++ b/src/main/java/org/opensearch/flowframework/rest/RestProvisionWorkflowAction.java @@ -22,7 +22,7 @@ import org.opensearch.flowframework.exception.FlowFrameworkException; import org.opensearch.flowframework.transport.ProvisionWorkflowAction; import org.opensearch.flowframework.transport.WorkflowRequest; -import org.opensearch.flowframework.util.RestActionUtils; +import org.opensearch.flowframework.util.TenantAwareHelper; import org.opensearch.rest.BaseRestHandler; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; @@ -86,7 +86,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli RestStatus.FORBIDDEN ); } - String tenantId = RestActionUtils.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); + String tenantId = TenantAwareHelper.getTenantID(flowFrameworkFeatureEnabledSetting.isMultiTenancyEnabled(), request); // Validate params if (workflowId == null) { throw new FlowFrameworkException("workflow_id cannot be null", RestStatus.BAD_REQUEST); diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index 51e796e98..52a7d2251 100644 --- a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java +++ b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java @@ -76,7 +76,7 @@ public class EncryptorUtils { private static final String WRAPPING_ALGORITHM = "AES/GCM/NOPADDING"; // concurrent map can't have null as a key. This key is to support single tenancy - public static final String DEFAULT_TENANT_ID = "03000200-0400-0500-0006-000700080009"; + public static final String DEFAULT_TENANT_ID = ""; private final ClusterService clusterService; private final Client client; diff --git a/src/main/java/org/opensearch/flowframework/util/RestActionUtils.java b/src/main/java/org/opensearch/flowframework/util/RestActionUtils.java deleted file mode 100644 index af784021d..000000000 --- a/src/main/java/org/opensearch/flowframework/util/RestActionUtils.java +++ /dev/null @@ -1,54 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ -package org.opensearch.flowframework.util; - -import org.opensearch.core.rest.RestStatus; -import org.opensearch.flowframework.common.CommonValue; -import org.opensearch.flowframework.exception.FlowFrameworkException; -import org.opensearch.rest.RestRequest; - -import java.util.List; -import java.util.Map; - -/** - * Utility methods for Rest Handlers - */ -public class RestActionUtils { - - private RestActionUtils() {} - - /** - * Finds the tenant id in the REST Headers - * @param isMultiTenancyEnabled whether multitenancy is enabled - * @param restRequest the RestRequest - * @return The tenant ID from the headers or null if multitenancy is not enabled - */ - public static String getTenantID(Boolean isMultiTenancyEnabled, RestRequest restRequest) { - if (isMultiTenancyEnabled) { - Map> headers = restRequest.getHeaders(); - if (headers.containsKey(CommonValue.TENANT_ID_HEADER)) { - List tenantIdList = headers.get(CommonValue.TENANT_ID_HEADER); - if (tenantIdList != null && !tenantIdList.isEmpty()) { - String tenantId = tenantIdList.get(0); - if (tenantId != null) { - return tenantId; - } else { - throw new FlowFrameworkException("Tenant ID can't be null", RestStatus.FORBIDDEN); - } - } else { - throw new FlowFrameworkException("Tenant ID header is present but has no value", RestStatus.FORBIDDEN); - } - } else { - throw new FlowFrameworkException("Tenant ID header is missing", RestStatus.FORBIDDEN); - } - } else { - return null; - } - } -} diff --git a/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java b/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java index eab891be9..eb78bc02f 100644 --- a/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java +++ b/src/main/java/org/opensearch/flowframework/util/TenantAwareHelper.java @@ -10,8 +10,12 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; +import org.opensearch.flowframework.common.CommonValue; import org.opensearch.flowframework.exception.FlowFrameworkException; +import org.opensearch.rest.RestRequest; +import java.util.List; +import java.util.Map; import java.util.Objects; /** @@ -29,7 +33,7 @@ public class TenantAwareHelper { */ public static boolean validateTenantId(boolean isMultiTenancyEnabled, String tenantId, ActionListener listener) { if (isMultiTenancyEnabled && tenantId == null) { - listener.onFailure(new FlowFrameworkException("You don't have permission to access this resource", RestStatus.FORBIDDEN)); + listener.onFailure(new FlowFrameworkException("No permission to access this resource", RestStatus.FORBIDDEN)); return false; } else { return true; @@ -52,8 +56,34 @@ public static boolean validateTenantResource( ActionListener listener ) { if (isMultiTenancyEnabled && !Objects.equals(tenantIdFromRequest, tenantIdFromResource)) { - listener.onFailure(new FlowFrameworkException("You don't have permission to access this resource", RestStatus.FORBIDDEN)); + listener.onFailure(new FlowFrameworkException("No permission to access this resource", RestStatus.FORBIDDEN)); return false; } else return true; } + + /** + * Finds the tenant id in the REST Headers + * @param isMultiTenancyEnabled whether multitenancy is enabled + * @param restRequest the RestRequest + * @return The tenant ID from the headers or null if multitenancy is not enabled + */ + public static String getTenantID(Boolean isMultiTenancyEnabled, RestRequest restRequest) { + if (!isMultiTenancyEnabled) { + return null; + } + + Map> headers = restRequest.getHeaders(); + + List tenantIdList = headers.get(CommonValue.TENANT_ID_HEADER); + if (tenantIdList == null || tenantIdList.isEmpty()) { + throw new FlowFrameworkException("Tenant ID header is missing or has no value", RestStatus.FORBIDDEN); + } + + String tenantId = tenantIdList.get(0); + if (tenantId == null) { + throw new FlowFrameworkException("Tenant ID can't be null", RestStatus.FORBIDDEN); + } + + return tenantId; + } } diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java index cbb6a42bd..1b793c10c 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkTenantAwareRestTestCase.java @@ -55,8 +55,8 @@ public abstract class FlowFrameworkTenantAwareRestTestCase extends FlowFramework protected static final String EMPTY_CONTENT = "{}"; // REST Response error reasons - protected static final String MISSING_TENANT_REASON = "Tenant ID header is missing"; - protected static final String NO_PERMISSION_REASON = "You don't have permission to access this resource"; + protected static final String MISSING_TENANT_REASON = "Tenant ID header is missing or has no value"; + protected static final String NO_PERMISSION_REASON = "No permission to access this resource"; protected String tenantId = randomAlphaOfLength(5); protected String otherTenantId = randomAlphaOfLength(6); diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 0b43d6c54..a3528a404 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -810,9 +810,7 @@ public void testAllDefaultUseCasesCreation() throws Exception { } } - // TODO Re-enable - // This test is currently failing due to ML Commons MLPredictionTaskRequest serialization - public void disableTemporarilytestSemanticSearchWithLocalModelEndToEnd() throws Exception { + public void testSemanticSearchWithLocalModelEndToEnd() throws Exception { // Checking if plugins are part of the integration test cluster so we can continue with this test List plugins = catPlugins(); if (!plugins.contains("opensearch-knn") && !plugins.contains("opensearch-neural-search")) { diff --git a/src/test/java/org/opensearch/flowframework/util/RestActionUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/RestActionUtilsTests.java deleted file mode 100644 index 36ff2cc13..000000000 --- a/src/test/java/org/opensearch/flowframework/util/RestActionUtilsTests.java +++ /dev/null @@ -1,103 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ -package org.opensearch.flowframework.util; - -import org.opensearch.core.rest.RestStatus; -import org.opensearch.flowframework.common.CommonValue; -import org.opensearch.flowframework.exception.FlowFrameworkException; -import org.opensearch.rest.RestRequest; -import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.test.rest.FakeRestRequest; -import org.junit.Assert; - -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -public class RestActionUtilsTests extends OpenSearchTestCase { - - public void testGetTenantID() { - String tenantId = "test-tenant"; - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(tenantId)); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - String actualTenantID = RestActionUtils.getTenantID(Boolean.TRUE, restRequest); - Assert.assertEquals(tenantId, actualTenantID); - } - - public void testGetTenantID_NullTenantID() { - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(null)); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - try { - RestActionUtils.getTenantID(Boolean.TRUE, restRequest); - Assert.fail("Expected FlowFrameworkException"); - } catch (Exception e) { - Assert.assertTrue(e instanceof FlowFrameworkException); - Assert.assertEquals(RestStatus.FORBIDDEN, ((FlowFrameworkException) e).status()); - Assert.assertEquals("Tenant ID can't be null", e.getMessage()); - } - } - - public void testGetTenantID_NoMultiTenancy() { - String tenantId = "test-tenant"; - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(tenantId)); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - String tenantID = RestActionUtils.getTenantID(Boolean.FALSE, restRequest); - Assert.assertNull(tenantID); - } - - public void testGetTenantID_EmptyTenantIDList() { - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, Collections.emptyList()); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - FlowFrameworkException exception = expectThrows( - FlowFrameworkException.class, - () -> RestActionUtils.getTenantID(Boolean.TRUE, restRequest) - ); - assertEquals(RestStatus.FORBIDDEN, exception.status()); - assertEquals("Tenant ID header is present but has no value", exception.getMessage()); - } - - public void testGetTenantID_MissingTenantIDHeader() { - Map> headers = new HashMap<>(); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - FlowFrameworkException exception = expectThrows( - FlowFrameworkException.class, - () -> RestActionUtils.getTenantID(Boolean.TRUE, restRequest) - ); - assertEquals(RestStatus.FORBIDDEN, exception.status()); - assertEquals("Tenant ID header is missing", exception.getMessage()); - } - - public void testGetTenantID_MultipleValues() { - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, List.of("tenant1", "tenant2")); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - String actualTenantID = RestActionUtils.getTenantID(Boolean.TRUE, restRequest); - assertEquals("tenant1", actualTenantID); - } - - public void testGetTenantID_EmptyStringTenantID() { - Map> headers = new HashMap<>(); - headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList("")); - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); - - String actualTenantID = RestActionUtils.getTenantID(Boolean.TRUE, restRequest); - assertEquals("", actualTenantID); - } -} diff --git a/src/test/java/org/opensearch/flowframework/util/TenantAwareHelperTests.java b/src/test/java/org/opensearch/flowframework/util/TenantAwareHelperTests.java index a6fbc7d1c..890e55d4c 100644 --- a/src/test/java/org/opensearch/flowframework/util/TenantAwareHelperTests.java +++ b/src/test/java/org/opensearch/flowframework/util/TenantAwareHelperTests.java @@ -10,24 +10,34 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.core.rest.RestStatus; +import org.opensearch.flowframework.common.CommonValue; import org.opensearch.flowframework.exception.FlowFrameworkException; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.rest.FakeRestRequest; +import org.junit.Assert; import org.junit.Before; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.verify; -public class TenantAwareHelperTests { +public class TenantAwareHelperTests extends OpenSearchTestCase { + private static final String NO_PERMISSION_TO_ACCESS = "No permission to access this resource"; @Mock private ActionListener actionListener; @Before - public void setUp() { + public void setUp() throws Exception { + super.setUp(); MockitoAnnotations.openMocks(this); } @@ -37,8 +47,8 @@ public void testValidateTenantId_MultiTenancyEnabled_TenantIdNull() { ArgumentCaptor captor = ArgumentCaptor.forClass(FlowFrameworkException.class); verify(actionListener).onFailure(captor.capture()); FlowFrameworkException exception = captor.getValue(); - assert exception.status() == RestStatus.FORBIDDEN; - assert exception.getMessage().equals("You don't have permission to access this resource"); + assertEquals(RestStatus.FORBIDDEN, exception.status()); + assertEquals(NO_PERMISSION_TO_ACCESS, exception.getMessage()); } public void testValidateTenantId_MultiTenancyEnabled_TenantIdPresent() { @@ -55,8 +65,8 @@ public void testValidateTenantResource_MultiTenancyEnabled_TenantIdMismatch() { ArgumentCaptor captor = ArgumentCaptor.forClass(FlowFrameworkException.class); verify(actionListener).onFailure(captor.capture()); FlowFrameworkException exception = captor.getValue(); - assert exception.status() == RestStatus.FORBIDDEN; - assert exception.getMessage().equals("You don't have permission to access this resource"); + assertEquals(RestStatus.FORBIDDEN, exception.status()); + assertEquals(NO_PERMISSION_TO_ACCESS, exception.getMessage()); } public void testValidateTenantResource_MultiTenancyEnabled_TenantIdMatch() { @@ -66,4 +76,83 @@ public void testValidateTenantResource_MultiTenancyEnabled_TenantIdMatch() { public void testValidateTenantResource_MultiTenancyDisabled() { assertTrue(TenantAwareHelper.validateTenantResource(false, "_tenant_id", "different_tenant_id", actionListener)); } + + public void testGetTenantID() { + String tenantId = "test-tenant"; + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(tenantId)); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + String actualTenantID = TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest); + assertEquals(tenantId, actualTenantID); + } + + public void testGetTenantID_NullTenantID() { + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(null)); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + try { + TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest); + Assert.fail("Expected FlowFrameworkException"); + } catch (Exception e) { + assertTrue(e instanceof FlowFrameworkException); + assertEquals(RestStatus.FORBIDDEN, ((FlowFrameworkException) e).status()); + assertEquals("Tenant ID can't be null", e.getMessage()); + } + } + + public void testGetTenantID_NoMultiTenancy() { + String tenantId = "test-tenant"; + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList(tenantId)); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + String tenantID = TenantAwareHelper.getTenantID(Boolean.FALSE, restRequest); + assertNull(tenantID); + } + + public void testGetTenantID_EmptyTenantIDList() { + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, Collections.emptyList()); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + FlowFrameworkException exception = expectThrows( + FlowFrameworkException.class, + () -> TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest) + ); + assertEquals(RestStatus.FORBIDDEN, exception.status()); + assertEquals("Tenant ID header is missing or has no value", exception.getMessage()); + } + + public void testGetTenantID_MissingTenantIDHeader() { + Map> headers = new HashMap<>(); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + FlowFrameworkException exception = expectThrows( + FlowFrameworkException.class, + () -> TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest) + ); + assertEquals(RestStatus.FORBIDDEN, exception.status()); + assertEquals("Tenant ID header is missing or has no value", exception.getMessage()); + } + + public void testGetTenantID_MultipleValues() { + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, List.of("tenant1", "tenant2")); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + String actualTenantID = TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest); + assertEquals("tenant1", actualTenantID); + } + + public void testGetTenantID_EmptyStringTenantID() { + Map> headers = new HashMap<>(); + headers.put(CommonValue.TENANT_ID_HEADER, Collections.singletonList("")); + RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withHeaders(headers).build(); + + String actualTenantID = TenantAwareHelper.getTenantID(Boolean.TRUE, restRequest); + assertEquals("", actualTenantID); + } + }