test + cleanup

Author: slisaasquatch

src/main/java/com/saasquatch/jsonschemainferrer/JsonSchemaInferrer.java

@@ -16,6 +16,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
@@ -248,6 +249,16 @@ private Set<ObjectNode> processPrimitives(@Nonnull Collection<ValueNode> valueNo
     return anyOfs;
   }
 
+  private Stream<ObjectNode> handleEnumExtractionResults(
+      @Nonnull Set<Set<? extends JsonNode>> enumExtractionResults) {
+    return enumExtractionResults.stream().map(enumExtractionResult -> {
+      final ObjectNode schema = newObject();
+      schema.set(Consts.Fields.ENUM, newArray(enumExtractionResult));
+      processGenericSchemaFeature(schema, enumExtractionResult, null);
+      return schema;
+    });
+  }
+
   /**
    * Build {@code anyOf} from sample JSONs. Note that all the arrays and objects will be combined.
    *
@@ -283,6 +294,7 @@ private Set<ObjectNode> getAnyOfsFromSamples(
       processGenericSchemaFeature(anyOf, enumExtractionResult, null);
       anyOfs.add(anyOf);
     }
+    handleEnumExtractionResults(enumExtractionResults).forEach(anyOfs::add);
     // Objects
     final ObjectNode resultForObjects = processObjects(objectNodes);
     if (resultForObjects != null) {

src/test/java/com/saasquatch/jsonschemainferrer/JsonSchemaInferrerOptionsTest.java

@@ -9,10 +9,12 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.time.DayOfWeek;
 import java.time.Instant;
+import java.time.Month;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
@@ -531,6 +533,20 @@ public void testEnum() {
       assertEquals(ImmutableSet.of(jnf.textNode("foo"), jnf.numberNode(123)),
           stream(schema.get("enum")).collect(Collectors.toSet()));
     }
+    {
+      final JsonSchemaInferrer inferrer = JsonSchemaInferrer.newBuilder()
+          .setEnumExtractor(EnumExtractors.chained(EnumExtractors.validEnum(Month.class),
+              EnumExtractors.validEnum(DayOfWeek.class)))
+          .build();
+      final ObjectNode schema =
+          inferrer.inferForSamples(Arrays.asList(jnf.textNode("TUESDAY"), jnf.textNode("MARCH")));
+      final JsonNode anyOf = schema.get("anyOf");
+      assertTrue(anyOf.isArray());
+      assertTrue(
+          stream(anyOf).anyMatch(_anyOf -> _anyOf.path("enum").get(0).textValue().equals("MARCH")));
+      assertTrue(stream(anyOf)
+          .anyMatch(_anyOf -> _anyOf.path("enum").get(0).textValue().equals("TUESDAY")));
+    }
   }
 
   @Test