From 499572dad038ee00690e16135d9da5e6ae8e61e3 Mon Sep 17 00:00:00 2001
From: Alan Jowett <alan.jowett@microsoft.com>
Date: Fri, 25 Oct 2024 08:35:12 -0700
Subject: [PATCH] PR feedback

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>
---
 src/test/ebpf_yaml.cpp | 27 ++++++++++++++++-----------
 src/test/ebpf_yaml.hpp |  2 +-
 test-data/add.yaml     |  6 ++++++
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/test/ebpf_yaml.cpp b/src/test/ebpf_yaml.cpp
index 4dbb3acd3..882c81fc8 100644
--- a/src/test/ebpf_yaml.cpp
+++ b/src/test/ebpf_yaml.cpp
@@ -131,12 +131,14 @@ static std::map<std::string, std::set<std::string>> parse_invariants_to_check(co
     if (!case_node["invariants-to-check"].IsDefined() || case_node["invariants-to-check"].IsNull()) {
         return {};
     }
-
     std::map<std::string, std::set<std::string>> res;
-    for (const auto& node : case_node["invariants-to-check"]) {
-        res.emplace(node.first.as<string>(), vector_to_set(node.second.as<vector<string>>()));
+    try {
+        for (const auto& node : case_node["invariants-to-check"]) {
+            res.emplace(node.first.as<string>(), vector_to_set(node.second.as<vector<string>>()));
+        }
+    } catch (const YAML::Exception& e) {
+        throw std::runtime_error("Error parsing 'invariants-to-check': " + std::string(e.what()));
     }
-
     return res;
 }
 
@@ -278,14 +280,17 @@ std::optional<Failure> run_yaml_test_case(TestCase test_case, bool debug) {
     std::set<string> actual_messages = extract_messages(ss.str());
 
     for (auto& [label, expected_invariant] : test_case.invariants_to_check) {
-        ss.str("");
-        ss.clear();
-        if (!ebpf_check_constraints_at_label(ss, label, expected_invariant)) {
-            // If debug is enabled, print the output of ebpf_check_constraints_at_label.
-            if (debug) {
-                std::cout << ss.str();
+        ss = {};
+        try {
+            if (!ebpf_check_constraints_at_label(ss, label, expected_invariant)) {
+                // If debug is enabled, print the output of ebpf_check_constraints_at_label.
+                if (debug) {
+                    std::cout << ss.str();
+                }
+                actual_messages.insert(label + ": Concrete invariants do not match abstract invariants");
             }
-            actual_messages.insert(label + ": Concrete invariants do not match abstract invariants");
+        } catch (const std::exception& e) {
+            actual_messages.insert(label + ": Exception occurred during invariant check - " + std::string(e.what()));
         }
     }
 
diff --git a/src/test/ebpf_yaml.hpp b/src/test/ebpf_yaml.hpp
index cd336f6d3..8847d1197 100644
--- a/src/test/ebpf_yaml.hpp
+++ b/src/test/ebpf_yaml.hpp
@@ -14,7 +14,7 @@ struct TestCase {
     InstructionSeq instruction_seq;
     string_invariant expected_post_invariant;
     std::set<std::string> expected_messages;
-    std::map<std::string, std::set<std::string>> invariants_to_check; ///< Map from label to expected invariants, used for checking invariants at specific labels.
+    std::map<std::string, std::set<std::string>> invariants_to_check; ///< Map from label to expected invariants, used for validating program state at specific labels during test execution.
 };
 
 void foreach_suite(const std::string& path, const std::function<void(const TestCase&)>& f);
diff --git a/test-data/add.yaml b/test-data/add.yaml
index 790f271ef..87529985a 100644
--- a/test-data/add.yaml
+++ b/test-data/add.yaml
@@ -193,14 +193,17 @@ invariants-to-check:
     - r1.type=number
     - r1.svalue=1
     - r1.uvalue=1
+
   2: # After the second assignment, r2 should be 2.
     - r2.type=number
     - r2.svalue=2
     - r2.uvalue=2
+
   3: # After the addition, r1 should be 3.
     - r1.type=number
     - r1.svalue=3
     - r1.uvalue=3
+
   -2: # After the last assignment, r0 should be 3.
     - r0.type=number
     - r0.svalue=3
@@ -233,14 +236,17 @@ invariants-to-check:
     - r1.type=number
     - r1.svalue=1
     - r1.uvalue=1
+
   2: # After the second assignment, r2 should be 2.
     - r2.type=number
     - r2.svalue=2
     - r2.uvalue=2
+
   3: # After the addition, r1 should be 3. Invalid invariant to test negative case.
     - r1.type=number
     - r1.svalue=2
     - r1.uvalue=2
+
   -2: # After the last assignment, r0 should be 3.
     - r0.type=number
     - r0.svalue=3