Add TrainedModelDefinitionDocTests and persist missing field

elastic · Apr 7, 2021 · b8e55f8 · b8e55f8
1 parent b002dd1
commit b8e55f8
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/...main/java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelDefinitionDoc.java b/...main/java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelDefinitionDoc.java
@@ -44,6 +44,7 @@ private static ObjectParser<TrainedModelDefinitionDoc.Builder, Void> createParse
         ObjectParser<TrainedModelDefinitionDoc.Builder, Void> parser = new ObjectParser<>(NAME,
             ignoreUnknownFields,
             TrainedModelDefinitionDoc.Builder::new);
+        parser.declareString((a, b) -> {}, InferenceIndexConstants.DOC_TYPE);  // type is hard coded but must be parsed
         parser.declareString(TrainedModelDefinitionDoc.Builder::setModelId, TrainedModelConfig.MODEL_ID);
         parser.declareString(TrainedModelDefinitionDoc.Builder::setCompressedString, DEFINITION);
         parser.declareInt(TrainedModelDefinitionDoc.Builder::setDocNum, DOC_NUM);
@@ -136,6 +137,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(TrainedModelConfig.MODEL_ID.getPreferredName(), modelId);
         builder.field(DOC_NUM.getPreferredName(), docNum);
         builder.field(DEFINITION_LENGTH.getPreferredName(), definitionLength);
+        builder.field(TOTAL_DEFINITION_LENGTH.getPreferredName(), totalDefinitionLength);
         builder.field(COMPRESSION_VERSION.getPreferredName(), compressionVersion);
         builder.field(DEFINITION.getPreferredName(), compressedString);
         builder.field(EOS.getPreferredName(), eos);

diff --git a/...java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelDefinitionDocTests.java b/...java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelDefinitionDocTests.java
@@ -0,0 +1,42 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.ml.inference.persistence;
+
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.test.AbstractXContentTestCase;
+
+import java.io.IOException;
+
+public class TrainedModelDefinitionDocTests extends AbstractXContentTestCase<TrainedModelDefinitionDoc> {
+
+    private boolean isLenient = randomBoolean();
+
+    @Override
+    protected TrainedModelDefinitionDoc doParseInstance(XContentParser parser) throws IOException {
+        return TrainedModelDefinitionDoc.fromXContent(parser, isLenient).build();
+    }
+
+    @Override
+    protected boolean supportsUnknownFields() {
+        return isLenient;
+    }
+
+    @Override
+    protected TrainedModelDefinitionDoc createTestInstance() {
+        int length = randomIntBetween(1, 10);
+        return new TrainedModelDefinitionDoc.Builder()
+            .setModelId(randomAlphaOfLength(6))
+            .setDefinitionLength(length)
+            .setTotalDefinitionLength(randomIntBetween(length, length *2))
+            .setCompressedString(randomAlphaOfLength(length))
+            .setDocNum(randomIntBetween(0, 10))
+            .setCompressionVersion(randomIntBetween(1, 5))
+            .setEos(randomBoolean())
+            .build();
+    }
+}