Fix reconstruction of DynamicSchema from GrpcRootMessage

It failed in case nested GrpcSchemaMessage contained another nested messages or enums. Also, removed duplicate test.
leviysoft · Dec 21, 2023 · 3be390b · 3be390b
1 parent 5b05591
commit 3be390b
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 62 deletions.
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/grpc/GrpcExractor.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/grpc/GrpcExractor.scala
@@ -39,64 +39,52 @@ object GrpcExractor {
 
   def addSchemaToRegistry(schema: GrpcRootMessage, registry: DynamicSchema.Builder): Unit =
     schema match {
-      case GrpcMessageSchema(name, fields, oneofs, nested, nestedEnums) =>
-        val builder = MessageDefinition
-          .newBuilder(name)
-        fields.foreach {
-          case f if f.label == GrpcLabel.Optional =>
-            val oneOfBuilder = builder.addOneof(s"_${f.name}")
-            oneOfBuilder.addField(f.typeName, f.name, f.order)
-          case f =>
-            builder.addField(f.label.entryName, f.typeName, f.name, f.order)
-        }
-        oneofs.getOrElse(List.empty).foreach { oneof =>
-          val oneOfBuilder = builder.addOneof(oneof.name)
-          oneof.options.foreach { of =>
-            oneOfBuilder.addField(of.typeName, of.name, of.order)
-          }
-        }
-        nested.getOrElse(List.empty).foreach { nst =>
-          val nestedBuilder = MessageDefinition.newBuilder(nst.name)
-          nst.fields.foreach {
-            case f if f.label == GrpcLabel.Optional =>
-              val oneOfBuilder = nestedBuilder.addOneof(s"_${f.name}")
-              oneOfBuilder.addField(f.typeName, f.name, f.order)
-            case f =>
-              nestedBuilder.addField(f.label.entryName, f.typeName, f.name, f.order)
-          }
-          builder.addMessageDefinition(nestedBuilder.build())
-        }
-        nestedEnums.getOrElse(List.empty).foreach { nst =>
-          val nestedBuilder = EnumDefinition.newBuilder(nst.name)
-          nst.values.foreach { case (name, number) =>
-            ignore(
-              nestedBuilder
-                .addValue(
-                  name.asString,
-                  number.asInt
-                )
-            )
-          }
-          builder.addEnumDefinition(nestedBuilder.build())
-        }
+      case m: GrpcMessageSchema =>
+        ignore(registry.addMessageDefinition(buildMessageDefinition(m)))
+      case m: GrpcEnumSchema =>
+        ignore(registry.addEnumDefinition(buildEnumDefinition(m)))
+    }
 
-        ignore(registry.addMessageDefinition(builder.build()))
-      case GrpcEnumSchema(name, values) =>
-        val builder = EnumDefinition
-          .newBuilder(name)
-        values.foreach { case (name, number) =>
-          ignore(
-            builder
-              .addValue(
-                name.asString,
-                number.asInt
-              )
-          )
-        }
-        val enumDefinition = builder.build()
-        ignore(registry.addEnumDefinition(enumDefinition))
+  def buildMessageDefinition(gm: GrpcMessageSchema): MessageDefinition = {
+    val builder = MessageDefinition.newBuilder(gm.name)
+
+    gm.fields.foreach {
+      case f if f.label == GrpcLabel.Optional =>
+        val oneOfBuilder = builder.addOneof(s"_${f.name}")
+        oneOfBuilder.addField(f.typeName, f.name, f.order)
+      case f =>
+        builder.addField(f.label.entryName, f.typeName, f.name, f.order)
     }
 
+    gm.oneofs.getOrElse(List.empty).foreach { oneof =>
+      val oneOfBuilder = builder.addOneof(oneof.name)
+      oneof.options.foreach { of =>
+        oneOfBuilder.addField(of.typeName, of.name, of.order)
+      }
+    }
+
+    gm.nested
+      .getOrElse(List.empty)
+      .foreach(
+        buildMessageDefinition andThen builder.addMessageDefinition
+      )
+
+    gm.nestedEnums
+      .getOrElse(List.empty)
+      .foreach(
+        buildEnumDefinition andThen builder.addEnumDefinition
+      )
+
+    builder.build()
+  }
+
+  def buildEnumDefinition(ge: GrpcEnumSchema): EnumDefinition =
+    ge.values
+      .foldLeft(EnumDefinition.newBuilder(ge.name)) { case (builder, (name, number)) =>
+        builder.addValue(name.asString, number.asInt)
+      }
+      .build()
+
   implicit class FromGrpcProtoDefinition(private val definition: GrpcProtoDefinition) extends AnyVal {
     def toDynamicSchema: DynamicSchema = {
       val registryBuilder: DynamicSchema.Builder = DynamicSchema.newBuilder()

diff --git a/backend/mockingbird/src/test/resources/nested.proto b/backend/mockingbird/src/test/resources/nested.proto
@@ -23,6 +23,23 @@ message GetStocksResponse {
   message Stocks {
     repeated Stock stocks = 1;
   }
+  message Event {
+    enum Code {
+      C_OK = 0;
+      C_ERROR = 1;
+    }
+    message Data {
+      string value = 1;
+    }
+    message Error {
+      string info = 1;
+    }
+    Code code = 1;
+    oneof payload {
+      Data data = 4;
+      Error error = 100;
+    }
+  }
 }
 
 service StockService {

diff --git a/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/protobuf/MappersSpec.scala b/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/protobuf/MappersSpec.scala
@@ -30,6 +30,10 @@ object MappersSpec extends ZIOSpecDefault {
     ".GetStocksResponse.StockKinds",
     ".GetStocksResponse.Stock",
     ".GetStocksResponse.Stocks",
+    ".GetStocksResponse.Event",
+    ".GetStocksResponse.Event.Code",
+    ".GetStocksResponse.Event.Data",
+    ".GetStocksResponse.Event.Error",
   )
 
   override def spec: Spec[TestEnvironment & Scope, Any] =
@@ -49,13 +53,6 @@ object MappersSpec extends ZIOSpecDefault {
           protoDefinitionAgain = protoDefinition.toDynamicSchema.toGrpcProtoDefinition
         } yield assertTrue(protoDefinition == protoDefinitionAgain)
       },
-      test("Mappers from nested DynamicSchema to GrpcProtoDefinition and back are consistent") {
-        for {
-          content <- Utils.getProtoDescriptionFromResource("nested.proto")
-          schema          = DynamicSchema.parseFrom(content)
-          protoDefinition = schema.toGrpcProtoDefinition
-        } yield assertTrue(getAllTypes(protoDefinition) == allTypesInNested)
-      },
       test("Mappers from nested DynamicSchema to GrpcProtoDefinition and back are consistent") {
         for {
           content <- Utils.getProtoDescriptionFromResource("nested.proto")