Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify ml license checking with XpackLicenseState internals #52684

Merged
merged 2 commits into from
Feb 26, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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> {

Expand Down Expand Up @@ -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));
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The licensing for inference models is not straightforward as different models can be licensed at different levels can we keep these assertions please so we know exactly what the behaviour is. The test that this code calls another piece of code is not sufficient. And for the gold tests below please

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These assertions were removed because the method to be tested was reduced to just a single line which simply makes a call to XPackLicenseState#isAllowedByLicense(licenseLevel). Since tests use mock object, there really was no other logic to be tested for this method.

But if we add some of the original logic back, e.g. the platinum license check as discussed above, the relevant tests are then become necessary again and need to be added back as well.

Copy link
Member

@davidkyle davidkyle Feb 24, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know it's a one liner and the mocking makes the test awkward but it may not always be a one-liner if someone refactors the function and I'd prefer to have the safety net of asserting the expected behaviour.

That said, the best kind of code is deleted code so why not remove the method completely and call XPackLicenseState#isAllowedByLicense() directly then the tests can go.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a great suggestion. I updated the code accordingly.

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));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down