From 2731ee43ca44c521eb813cd05b2f715772c58b8b Mon Sep 17 00:00:00 2001
From: Rohit Krishnan <rohitkrishnan@google.com>
Date: Mon, 18 Sep 2023 16:06:06 -0400
Subject: [PATCH] Fix bug resolving enum across multiple files.

---
 src/utils.rs              | 35 +++++++++++++++++++++++++----------
 tests/enum.proto          |  5 +++++
 tests/enum_imported.proto | 10 ++++++++++
 tests/enum_tests.rs       |  9 +++------
 tests/simple.proto        |  2 ++
 5 files changed, 45 insertions(+), 16 deletions(-)
 create mode 100644 tests/enum_imported.proto

diff --git a/src/utils.rs b/src/utils.rs
index 204fd70..b135a70 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -216,9 +216,19 @@ pub fn find_enum_value_by_name(
   let package_names = enum_name.split('.').filter(|v| !v.is_empty()).collect::<Vec<_>>();
   if let Some((_name, package)) = package_names.split_last() {
     let package = package.join(".");
-    descriptors.values()
-      .find(|fd| fd.package.clone().unwrap_or_default() == package)
-      .and_then(|fd| find_enum_value_by_name_in_message(&fd.enum_type, enum_name, enum_value))
+    let result = descriptors.values()
+      .find_map(|fd| {
+        if fd.package.clone().unwrap_or_default() == package {
+          find_enum_value_by_name_in_message(&fd.enum_type, enum_name, enum_value)
+        } else {
+          None
+        }
+      });
+    if result.is_some() {
+      result
+    } else {
+      None
+    }
   } else {
     None
   }
@@ -235,15 +245,20 @@ pub fn find_enum_by_name(
   trace!("package_names={:?}", package_names);
   if let Some((_name, package)) = package_names.split_last() {
     let package = package.join(".");
-    descriptors.file.iter()
-      .find(|fd| {
+    let result = descriptors.file.iter()
+      .find_map(|fd| {
         if let Some(fd_package) = &fd.package {
-          package == fd_package.as_str()
-        } else {
-          false
+          if package == fd_package.as_str() {
+            return find_enum_by_name_in_message(&fd.enum_type, enum_name)  
+          }
         }
-      })
-      .and_then(|fd| find_enum_by_name_in_message(&fd.enum_type, enum_name))
+        return None
+      });
+      if result.is_some() {
+        result
+      } else {
+        None
+      }
   } else {
     None
   }
diff --git a/tests/enum.proto b/tests/enum.proto
index b3ec414..df0015d 100644
--- a/tests/enum.proto
+++ b/tests/enum.proto
@@ -1,5 +1,9 @@
 syntax = "proto3";
 
+package example.enum.package;
+
+import "enum_imported.proto";
+
 enum Values {
   NONE = 0;
   A = 1;
@@ -9,6 +13,7 @@ enum Values {
 
 message MessageIn {
   repeated Values in = 1;
+  repeated Values2 in2 = 2;
 }
 
 message MessageOut {
diff --git a/tests/enum_imported.proto b/tests/enum_imported.proto
new file mode 100644
index 0000000..717ee44
--- /dev/null
+++ b/tests/enum_imported.proto
@@ -0,0 +1,10 @@
+syntax = "proto3";
+
+package example.enum.package;
+
+enum Values2 {
+    VALUES_2_NONE = 0;
+    VALUES_2_A = 1;
+    VALUES_2_B = 2;
+    VALUES_2_C = 3;
+}
\ No newline at end of file
diff --git a/tests/enum_tests.rs b/tests/enum_tests.rs
index 23ac93a..a8d7550 100644
--- a/tests/enum_tests.rs
+++ b/tests/enum_tests.rs
@@ -27,6 +27,9 @@ async fn repeated_enum_test() {
 
         "in": [
           "matching(equalTo, 'A')"
+        ],
+        "in2": [
+          "matching(equalTo, 'VALUES_2_A')"
         ]
       })).await;
       i
@@ -68,12 +71,6 @@ async fn repeated_enum_test() {
         wire_type: WireType::LengthDelimited,
         data: ProtobufFieldData::Enum(1, enum_descriptor.clone())
       },
-      // ProtobufField {
-      //   field_num: 1,
-      //   field_name: "in".to_string(),
-      //   wire_type: WireType::LengthDelimited,
-      //   data: ProtobufFieldData::Enum(1, enum_descriptor.clone())
-      // }
     ];
 
     let result = compare_message(
diff --git a/tests/simple.proto b/tests/simple.proto
index 25ed441..c1c95ef 100644
--- a/tests/simple.proto
+++ b/tests/simple.proto
@@ -1,5 +1,7 @@
 syntax = "proto3";
 
+package com.pact.protobuf.example;
+
 message MessageIn {
   bool in = 1;
 }