From ecf933dbd975ff3b8e54ca53fe89a20f91f2b202 Mon Sep 17 00:00:00 2001
From: Dimitris Koutsogiorgas <dimitris@squareup.com>
Date: Mon, 9 Oct 2023 10:31:27 +0100
Subject: [PATCH] WIP

---
 .../squareup/wire/schema/PartitionedSchema.kt |  2 +-
 .../wire/schema/ManifestPartitionTest.kt      | 56 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/PartitionedSchema.kt b/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/PartitionedSchema.kt
index 90ce712fcd..f3c2470425 100644
--- a/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/PartitionedSchema.kt
+++ b/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/PartitionedSchema.kt
@@ -113,5 +113,5 @@ internal fun Schema.partition(modules: Map<String, Module>): PartitionedSchema {
     }
   }
 
-  return PartitionedSchema(partitions, warnings, errors)
+    return PartitionedSchema(partitions, warnings, errors)
 }
diff --git a/wire-schema/src/jvmTest/kotlin/com/squareup/wire/schema/ManifestPartitionTest.kt b/wire-schema/src/jvmTest/kotlin/com/squareup/wire/schema/ManifestPartitionTest.kt
index 4edc75cae1..95c7733a4f 100644
--- a/wire-schema/src/jvmTest/kotlin/com/squareup/wire/schema/ManifestPartitionTest.kt
+++ b/wire-schema/src/jvmTest/kotlin/com/squareup/wire/schema/ManifestPartitionTest.kt
@@ -177,6 +177,62 @@ class ManifestPartitionTest {
       """.trimMargin(),
     )
   }
+
+  @Test fun something() {
+    val schema = buildSchema {
+      add(
+        "first.proto".toPath(),
+        """
+          |syntax = "proto2";
+          |package com.squareup.wire.first;
+          |
+          |message Card {
+          |  enum Brand {
+          |     DO_NOT_USE = 1;
+          |  }
+          |  optional string name = 1;
+          |}
+          |
+        """.trimMargin(),
+      )
+
+        add(
+          "second.proto".toPath(),
+          """
+          |syntax = "proto2";
+          |package com.squareup.wire.second;
+          |
+          |import "first.proto";
+          |
+          |message Details {
+          |  optional com.squareup.wire.first.Card.Brand brand = 1;
+          |}
+          |
+        """.trimMargin(),
+        )
+    }
+
+    val modules = mapOf(
+      "first" to WireRun.Module(
+        pruningRules = PruningRules.Builder()
+          .addRoot("com.squareup.wire.first.Card")
+          .build(),
+      ),
+      "second" to WireRun.Module(
+        dependencies = setOf("first"),
+        pruningRules = PruningRules.Builder()
+          .addRoot("com.squareup.wire.second.Details")
+          .build(),
+      ),
+    )
+
+    val partitionedSchema = schema.partition(modules)
+
+    val firstPartition = partitionedSchema.partitions.getValue("first")
+
+    assertThat(firstPartition.types).contains(ProtoType.get("com.squareup.wire.first.Card"))
+    assertThat(firstPartition.types).contains(ProtoType.get("com.squareup.wire.first.Card.Brand"))
+  }
 }
 
 private fun Schema.getMessage(name: String): MessageType {