Java (Pseudo) Union type with custom Jackson (de)serializer

VariantSerializationTest.java

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import lombok.Builder;
import lombok.Data;
import lombok.extern.jackson.Jacksonized;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;


@Slf4j
public class VariantSerializationTest {

    private static final Map<String, Object> RAW_EXECUTABLE =
            Map.of("cmd", "ls", "args", List.of("-la"));

    private static final Executable EXECUTABLE = Executable.builder()
            .cmd("ls")
            .args(List.of("-la"))
            .build();

    private static final ObjectMapper MAPPER = new ObjectMapper();

    static {
        // java.lang.Optional support
        MAPPER.registerModules(new Jdk8Module());
    }


    @Test
    public void noNullableVariant() throws Throwable {
        var obj = new ParentWithNullableUnion(null);
        var json = MAPPER.writeValueAsString(obj);
        var parsed = MAPPER.readValue(json, ParentWithNullableUnion.class);
        assertNull(parsed.union());
    }

    @Test
    public void nullOptionalVariant() throws Throwable {
        var obj = new ParentWithOptionalUnion(Optional.empty());
        var json = MAPPER.writeValueAsString(obj);
        var parsed = MAPPER.readValue(json, ParentWithOptionalUnion.class);
        assertTrue(parsed.union().isEmpty());
    }

    @Test
    public void optionalStringVariant() throws Throwable {
        var parsed = runOptionalParentTest(
                new StringVariant("ls"),
                StringVariant.class,
                "ls"
        );
        assertEquals("ls", parsed.getValue());
    }

    @Test
    public void optionalListVariant() throws Throwable {

        var parsed = runOptionalParentTest(
                new ListVariant(List.of(EXECUTABLE)),
                ListVariant.class,
                List.of(RAW_EXECUTABLE)
        );
        assertEquals(List.of(EXECUTABLE), parsed.getValue());
    }

    @Test
    public void optionalObjectVariant() throws Throwable {
        var parsed = runOptionalParentTest(
                new ObjectVariant(EXECUTABLE),
                ObjectVariant.class,
                RAW_EXECUTABLE
        );
        assertEquals(EXECUTABLE, parsed.getValue());
    }

    @Test
    public void testGetter() {
        MyUnion union = new ListVariant(List.of());
        assertTrue(union.is(ListVariant.class));
        assertFalse(union.is(StringVariant.class));
        assertFalse(union.is(ObjectVariant.class));
        assertNotNull(union.get(ListVariant.class));
        assertThrows(RuntimeException.class, () -> {
            union.get(ObjectVariant.class);
        });
        assertThrows(RuntimeException.class, () -> {
            union.get(StringVariant.class);
        });
    }

    private <T> T runOptionalParentTest(MyUnion union, Class<T> clz, Object raw) throws Throwable {
        var obj = new ParentWithOptionalUnion(Optional.of(union));
        // Serialize
        var json = MAPPER.writeValueAsString(obj);
        //Assert raw serialized value too
        var asMap = MAPPER.readValue(json, Map.class);
        assertEquals(asMap.get("union"), raw);
        // Deserialize
        var parsed = MAPPER.readValue(json, ParentWithOptionalUnion.class);
        var type = parsed.union().orElseThrow();
        assertTrue(clz.isInstance(type));
        return clz.cast(type);
    }
}

/**
 * Store union as Optional
 */
record ParentWithOptionalUnion(Optional<MyUnion> union) {

}

/**
 * Store nullable union
 */
record ParentWithNullableUnion(MyUnion union) {

}

/**
 * Union interface
 */
@JsonSerialize(using = UnionSerializer.class)
@JsonDeserialize(using = UnionDeserializer.class)
interface MyUnion {
    default <T extends MyUnion> T get(Class<T> clz) {
        return clz.cast(this);
    }

    default <T extends MyUnion> boolean is(Class<T> clz) {
        return clz.isInstance(this);
    }
}

@Builder
@Data
@Jacksonized
class Executable {
    private final String cmd;
    private final List<String> args;
}


@Builder
@Data
@Jacksonized
class StringVariant implements MyUnion {
    private final String value;
}


@Builder
@Data
@Jacksonized
class ListVariant implements MyUnion {
    private final List<Executable> value;
}


@Builder
@Data
@Jacksonized
class ObjectVariant implements MyUnion {
    private final Executable value;
}

class UnionSerializer extends JsonSerializer<MyUnion> {

    @Override
    public void serialize(
            MyUnion value, JsonGenerator gen, SerializerProvider provider)
            throws IOException, JsonProcessingException {

        var mapper = (ObjectMapper) gen.getCodec();
        if (value instanceof StringVariant a) {
            gen.writeString(a.getValue());
        } else if (value instanceof ListVariant a) {
            gen.writeRawValue(mapper.writeValueAsString(a.getValue()));
        } else if (value instanceof ObjectVariant a) {
            gen.writeRawValue(mapper.writeValueAsString(a.getValue()));
        } else {
            throw new IOException("Unsupported union type");
        }
    }
}

@Slf4j
class UnionDeserializer extends JsonDeserializer<MyUnion> {

    @Override
    public MyUnion deserialize(JsonParser parser, DeserializationContext ctxt)
            throws IOException, JsonProcessingException {

        var root = parser.getCodec().readTree(parser);
        var mapper = (ObjectMapper) parser.getCodec();

        if (root.isValueNode()) {
            return new StringVariant(mapper.readValue(root.toString(),
                    String.class));
        }
        // Must be before the `isContainerNode` check
        if (root.isObject()) {
            return new ObjectVariant(mapper.readValue(root.toString(),
                    Executable.class));
        }
        if (root.isContainerNode()) {
            return new ListVariant(mapper.readValue(root.toString(),
                    new TypeReference<List<Executable>>() {
                    })
            );
        }
        throw new IOException("Unable to deserialize union type");
    }
}