handle nested message types

Fadelis · Sep 2, 2021 · 66013df · 66013df
1 parent 0444f21
commit 66013df
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 7 deletions.
diff --git a/protoc-gen-java-optional-test/src/test/java/DataTypesTest.java b/protoc-gen-java-optional-test/src/test/java/DataTypesTest.java
@@ -4,6 +4,8 @@
 import org.junit.jupiter.api.Test;
 import org.test.datatypes.TestEnum;
 import org.test.datatypes.TestMessage;
+import org.test.datatypes.TestMessage.NestedTestMessage;
+import org.test.datatypes.TestMessage.NestedTestMessage.DeepNestedTestMessage;
 import org.test.datatypes.TestSubMessage;
 
 class DataTypesTest {
@@ -182,6 +184,39 @@ void should_have_empty_optional_on_no_data_in_builder() {
     assertThat(builder.getOptionalSubMessage()).isEmpty();
   }
 
+  @Test
+  void should_have_optional_methods_in_nested_types() {
+    NestedTestMessage nestedTestMessage = NestedTestMessage.getDefaultInstance();
+    DeepNestedTestMessage deepNestedTestMessage = DeepNestedTestMessage.getDefaultInstance();
+
+    assertThat(nestedTestMessage.getOptionalOptionalString()).isEmpty();
+    assertThat(nestedTestMessage.getOptionalSubMessage()).isEmpty();
+    assertThat(deepNestedTestMessage.getOptionalOptionalString()).isEmpty();
+    assertThat(deepNestedTestMessage.getOptionalSubMessage()).isEmpty();
+  }
+
+  @Test
+  void should_have_optional_builder_methods_in_nested_types() {
+    NestedTestMessage.Builder nestedTestMessageBuilder = NestedTestMessage.newBuilder();
+    DeepNestedTestMessage.Builder deepNestedTestMessageBuilder = DeepNestedTestMessage.newBuilder();
+
+    assertThat(nestedTestMessageBuilder.getOptionalOptionalString()).isEmpty();
+    assertThat(nestedTestMessageBuilder.getOptionalSubMessage()).isEmpty();
+    assertThat(deepNestedTestMessageBuilder.getOptionalOptionalString()).isEmpty();
+    assertThat(deepNestedTestMessageBuilder.getOptionalSubMessage()).isEmpty();
+  }
+
+  @Test
+  void should_have_set_or_clear_builder_methods_in_nested_types() {
+    NestedTestMessage.Builder nestedTestMessageBuilder = NestedTestMessage.newBuilder();
+    DeepNestedTestMessage.Builder deepNestedTestMessageBuilder = DeepNestedTestMessage.newBuilder();
+
+    assertThat(nestedTestMessageBuilder.setOrClearOptionalString(null)).isSameAs(nestedTestMessageBuilder);
+    assertThat(nestedTestMessageBuilder.setOrClearSubMessage(null)).isSameAs(nestedTestMessageBuilder);
+    assertThat(deepNestedTestMessageBuilder.setOrClearOptionalString(null)).isSameAs(deepNestedTestMessageBuilder);
+    assertThat(deepNestedTestMessageBuilder.setOrClearSubMessage(null)).isSameAs(deepNestedTestMessageBuilder);
+  }
+
   private static TestMessage.Builder populatedTestMessage() {
     return TestMessage.newBuilder()
         .setOneofString("one-of")

diff --git a/protoc-gen-java-optional-test/src/test/resources/data-types.proto b/protoc-gen-java-optional-test/src/test/resources/data-types.proto
@@ -21,10 +21,22 @@ message TestMessage {
     value2 = 2;
   }
 
+  message NestedTestMessage {
+    message DeepNestedTestMessage {
+      string string = 1;
+      optional string optional_string = 2;
+      TestSubMessage sub_message = 3;
+    }
+
+    string string = 1;
+    optional string optional_string = 2;
+    TestSubMessage sub_message = 3;
+  }
+
   string string = 1;
   int32 int32 = 2;
   bool bool = 3;
-  TestSubMessage subMessage = 5;
+  TestSubMessage sub_message = 5;
   double double = 6;
   float float = 7;
   bytes bytes = 8;
@@ -60,7 +72,7 @@ message RepeatedTestMessage {
   repeated string string = 1;
   repeated int32 int32 = 2;
   repeated bool bool = 3;
-  repeated TestSubMessage subMessage = 5;
+  repeated TestSubMessage sub_message = 5;
   repeated double double = 6;
   repeated float float = 7;
   repeated bytes bytes = 8;

diff --git a/protoc-gen-java-optional/src/main/java/org/grpcmock/protoc/plugin/OptionalGenerator.java b/protoc-gen-java-optional/src/main/java/org/grpcmock/protoc/plugin/OptionalGenerator.java
@@ -94,16 +94,29 @@ private Stream<File> handleProtoFile(FileDescriptorProto fileDescriptor) {
         : protoPackage;
 
     return fileDescriptor.getMessageTypeList().stream()
-        .flatMap(messageDescriptor -> handleMessage(messageDescriptor, protoPackage, javaPackage));
+        .flatMap(descriptor -> handleMessageProto(descriptor, descriptor.getName(), protoPackage, javaPackage));
   }
 
-  private Stream<File> handleMessage(DescriptorProto messageDescriptor, String protoPackage, String javaPackage) {
-    String fileName = javaPackage.replace(".", DIR_SEPARATOR) + DIR_SEPARATOR + messageDescriptor.getName() + JAVA_EXTENSION;
+  private Stream<File> handleMessageProto(
+      DescriptorProto messageDescriptor,
+      String fileName,
+      String protoPackage,
+      String javaPackage
+  ) {
+    String filePath = javaPackage.replace(".", DIR_SEPARATOR) + DIR_SEPARATOR + fileName + JAVA_EXTENSION;
     String fullMethodName = protoPackage + "." + messageDescriptor.getName();
 
+    return Stream.concat(
+        handleMessage(messageDescriptor, filePath, fullMethodName),
+        messageDescriptor.getNestedTypeList().stream()
+            .filter(nestedDescriptor -> !nestedDescriptor.getOptions().getMapEntry())
+            .flatMap(nestedDescriptor -> handleMessageProto(nestedDescriptor, fileName, fullMethodName, javaPackage)));
+  }
+
+  private Stream<File> handleMessage(DescriptorProto messageDescriptor, String filePath, String fullMethodName) {
     return Stream.of(
-        createFile(messageDescriptor, fileName, fullMethodName, BUILDER_SCOPE, this::createBuilderMethods),
-        createFile(messageDescriptor, fileName, fullMethodName, CLASS_SCOPE, this::createClassMethods))
+            createFile(messageDescriptor, filePath, fullMethodName, BUILDER_SCOPE, this::createBuilderMethods),
+            createFile(messageDescriptor, filePath, fullMethodName, CLASS_SCOPE, this::createClassMethods))
         .filter(Optional::isPresent)
         .map(Optional::get);
   }