Simplify ml license checking with XpackLicenseState internals (elasti…

…c#52684) This change removes TrainedModelConfig#isAvailableWithLicense method with calls to XPackLicenseState#isAllowedByLicense. Please note there are subtle changes to the code logic. But they are the right changes: * Instead of Platinum license, Enterprise license nows guarantees availability. * No explicit check when the license requirement is basic. Since basic license is always available, this check is unnecessary. * Trial license is always allowed.
ywangd · Feb 27, 2020 · f08772c · f08772c
1 parent f5c4e92
commit f08772c
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 85 deletions.
diff --git a/...ugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java b/...ugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java
@@ -21,7 +21,6 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.license.License;
-import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.xpack.core.common.time.TimeUtils;
 import org.elasticsearch.xpack.core.ml.inference.persistence.InferenceIndexConstants;
 import org.elasticsearch.xpack.core.ml.job.messages.Messages;
@@ -236,21 +235,6 @@ public License.OperationMode getLicenseLevel() {
         return licenseLevel;
     }
 
-    public boolean isAvailableWithLicense(XPackLicenseState licenseState) {
-        // Basic is always true
-        if (licenseLevel.equals(License.OperationMode.BASIC)) {
-            return true;
-        }
-
-        // The model license does not matter, Platinum license gets the same functions as the highest license
-        if (licenseState.isAllowedByLicense(License.OperationMode.PLATINUM)) {
-            return true;
-        }
-
-        // catch the rest, if the license is active and is at least the required model license
-        return licenseState.isAllowedByLicense(licenseLevel, true, false);
-    }
-
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(modelId);

diff --git a/...core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java b/...core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java
@@ -20,7 +20,6 @@
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.license.License;
-import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.search.SearchModule;
 import org.elasticsearch.test.AbstractSerializingTestCase;
 import org.elasticsearch.xpack.core.ml.job.messages.Messages;
@@ -44,10 +43,6 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.not;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyBoolean;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 public class TrainedModelConfigTests extends AbstractSerializingTestCase<TrainedModelConfig> {
 
@@ -306,66 +301,4 @@ public void testSerializationWithCompressedLazyDefinition() throws IOException {
             .assertToXContentEquivalence(true)
             .test();
     }
-
-    public void testIsAvailableWithLicense() {
-        TrainedModelConfig.Builder builder = createTestInstance(randomAlphaOfLength(10));
-        XPackLicenseState licenseState = mock(XPackLicenseState.class);
-
-        // Reject everything
-        when(licenseState.isAllowedByLicense(any(License.OperationMode.class), anyBoolean(), anyBoolean())).thenAnswer(
-            invocationOnMock -> {
-                final Object[] arguments = invocationOnMock.getArguments();
-                assertTrue((boolean) arguments[1]); // ensure the call is made to require active license
-                return false;
-            }
-        );
-        assertFalse(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState));
-        assertFalse(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState));
-        assertFalse(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState));
-        // Basic license always works not matter what
-        assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState));
-    }
-
-    public void testActivePlatinumLicenseAlwaysWorks() {
-        TrainedModelConfig.Builder builder = createTestInstance(randomAlphaOfLength(10));
-        XPackLicenseState licenseState = mock(XPackLicenseState.class);
-
-        when(licenseState.isAllowedByLicense(License.OperationMode.PLATINUM)).thenReturn(true);
-
-        // Active Platinum license functions the same as Enterprise license (highest) and should always work
-        when(licenseState.isAllowedByLicense(any(License.OperationMode.class), anyBoolean(), anyBoolean())).thenAnswer(
-            invocationOnMock -> {
-                final Object[] arguments = invocationOnMock.getArguments();
-                assertEquals(License.OperationMode.PLATINUM, arguments[0]);
-                assertTrue((boolean) arguments[1]); // ensure the call is made to require active license
-                assertTrue((boolean) arguments[2]);
-                return true;
-            }
-        );
-        assertTrue(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState));
-        assertTrue(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState));
-        assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState));
-        assertTrue(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState));
-    }
-
-    public void testActiveGoldLicenseWillWorkWhenRequiredLevelIsGold() {
-        TrainedModelConfig.Builder builder = createTestInstance(randomAlphaOfLength(10));
-        XPackLicenseState licenseState = mock(XPackLicenseState.class);
-
-        // Active Gold license should work when required level is gold
-        when(licenseState.isAllowedByLicense(any(License.OperationMode.class), anyBoolean(), anyBoolean())).thenAnswer(
-            invocationOnMock -> {
-                final Object[] arguments = invocationOnMock.getArguments();
-                assertTrue((boolean) arguments[1]); // ensure the call is made to require active license
-                if (License.OperationMode.PLATINUM == arguments[0] && Boolean.TRUE.equals(arguments[2])) {
-                    return false;
-                } else
-                    return License.OperationMode.GOLD == arguments[0] && Boolean.FALSE.equals(arguments[2]);
-            }
-        );
-        assertFalse(builder.setLicenseLevel(License.OperationMode.ENTERPRISE.description()).build().isAvailableWithLicense(licenseState));
-        assertFalse(builder.setLicenseLevel(License.OperationMode.PLATINUM.description()).build().isAvailableWithLicense(licenseState));
-        assertTrue(builder.setLicenseLevel(License.OperationMode.BASIC.description()).build().isAvailableWithLicense(licenseState));
-        assertTrue(builder.setLicenseLevel(License.OperationMode.GOLD.description()).build().isAvailableWithLicense(licenseState));
-    }
 }
diff --git a/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java
@@ -79,8 +79,8 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
         } else {
             trainedModelProvider.getTrainedModel(request.getModelId(), false, ActionListener.wrap(
                 trainedModelConfig -> {
-                    responseBuilder.setLicensed(trainedModelConfig.isAvailableWithLicense(licenseState));
-                    if (trainedModelConfig.isAvailableWithLicense(licenseState) || request.isPreviouslyLicensed()) {
+                    responseBuilder.setLicensed(licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel()));
+                    if (licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel()) || request.isPreviouslyLicensed()) {
                         this.modelLoadingService.getModel(request.getModelId(), getModelListener);
                     } else {
                         listener.onFailure(LicenseUtils.newComplianceException(XPackField.MACHINE_LEARNING));