From 531dbed7cdeb7061368bce7ced95ffebc505ec4e Mon Sep 17 00:00:00 2001
From: Jack Khuu <jackkhuu@meta.com>
Date: Thu, 30 Nov 2023 15:18:24 -0800
Subject: [PATCH] InspectorUtil Test: find_populated_event (#1326)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/1326

Add missing test for InspectorUtil

More to come, adding these iteratively quick reviews

Reviewed By: Olivia-liu

Differential Revision: D51610165

fbshipit-source-id: cc0f2b4de6266e0220517956ee1fbdd7a572123a
---
 sdk/inspector/tests/TARGETS                 |  1 +
 sdk/inspector/tests/inspector_utils_test.py | 57 +++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/sdk/inspector/tests/TARGETS b/sdk/inspector/tests/TARGETS
index c8a49d7e4f..0e6d06e776 100644
--- a/sdk/inspector/tests/TARGETS
+++ b/sdk/inspector/tests/TARGETS
@@ -32,6 +32,7 @@ python_unittest(
         "//executorch/sdk:lib",
         "//executorch/sdk/debug_format:base_schema",
         "//executorch/sdk/debug_format:et_schema",
+        "//executorch/sdk/etdump:schema_flatcc",
         "//executorch/sdk/etrecord/tests:etrecord_test_library",
         "//executorch/sdk/inspector:inspector_utils",
     ],
diff --git a/sdk/inspector/tests/inspector_utils_test.py b/sdk/inspector/tests/inspector_utils_test.py
index a5b9095353..a9a4ec7d52 100644
--- a/sdk/inspector/tests/inspector_utils_test.py
+++ b/sdk/inspector/tests/inspector_utils_test.py
@@ -17,11 +17,13 @@
 )
 
 from executorch.sdk.debug_format.et_schema import FXOperatorGraph
+from executorch.sdk.etdump import schema_flatcc as flatcc
 
 from executorch.sdk.etrecord.tests.etrecord_test import TestETRecord
 from executorch.sdk.inspector._inspector_utils import (
     create_debug_handle_to_op_node_mapping,
     EDGE_DIALECT_GRAPH_KEY,
+    find_populated_event,
     gen_graphs_from_etrecord,
 )
 
@@ -57,6 +59,61 @@ def test_create_debug_handle_to_op_node_mapping(self):
 
         self.assertEqual(debug_handle_to_op_node_map, expected_mapping)
 
+    def test_find_populated_event(self):
+        profile_event = flatcc.ProfileEvent(
+            name="test_profile_event",
+            chain_index=1,
+            instruction_id=1,
+            delegate_debug_id_str="",
+            delegate_debug_id_int=-1,
+            delegate_debug_metadata="",
+            start_time=1001,
+            end_time=2002,
+        )
+        debug_event = flatcc.DebugEvent(
+            chain_index=1,
+            instruction_id=0,
+            debug_entry=flatcc.Value(
+                val=flatcc.ValueType.TENSOR.value,
+                tensor=flatcc.Tensor(
+                    scalar_type=flatcc.ScalarType.INT,
+                    sizes=[1],
+                    strides=[1],
+                    offset=12345,
+                ),
+                int_value=flatcc.Int(1),
+                float_value=flatcc.Float(1.0),
+                double_value=flatcc.Double(1.0),
+                bool_value=flatcc.Bool(False),
+                output=flatcc.Bool(True),
+            ),
+        )
+
+        # Profile Populated
+        event = flatcc.Event(
+            profile_event=profile_event, debug_event=None, allocation_event=None
+        )
+        self.assertEqual(find_populated_event(event), profile_event)
+
+        # Debug Populated
+        event = flatcc.Event(
+            profile_event=None, debug_event=debug_event, allocation_event=None
+        )
+        self.assertEqual(find_populated_event(event), debug_event)
+
+        # Neither Populated
+        event = flatcc.Event(
+            profile_event=None, debug_event=None, allocation_event=None
+        )
+        with self.assertRaises(ValueError):
+            self.assertEqual(find_populated_event(event), profile_event)
+
+        # Both Populated (Returns Profile Event)
+        event = flatcc.Event(
+            profile_event=profile_event, debug_event=debug_event, allocation_event=None
+        )
+        self.assertEqual(find_populated_event(event), profile_event)
+
 
 def gen_mock_operator_graph_with_expected_map() -> Tuple[
     OperatorGraph, Dict[int, OperatorNode]