From 039b1daa330e8993fdd2385dffccf1c84cfa3b49 Mon Sep 17 00:00:00 2001
From: Jose Garcia Crosta <jgcrosta@gmail.com>
Date: Fri, 30 Aug 2024 10:30:21 -0300
Subject: [PATCH] Add removed tests

---
 .../vulnerable-example/src/lib.rs                | 16 ++++++++++++++++
 .../vulnerable-example/src/lib.rs                | 16 ++++++++++++++++
 .../vulnerable-example/src/lib.rs                | 16 ++++++++++++++++
 .../vulnerable-example/src/lib.rs                | 16 ++++++++++++++++
 .../vulnerable-example/src/lib.rs                | 16 ++++++++++++++++
 5 files changed, 80 insertions(+)

diff --git a/test-cases/unsafe-expect/unsafe-expect-1/vulnerable-example/src/lib.rs b/test-cases/unsafe-expect/unsafe-expect-1/vulnerable-example/src/lib.rs
index 064ceb62..40d5cd8b 100644
--- a/test-cases/unsafe-expect/unsafe-expect-1/vulnerable-example/src/lib.rs
+++ b/test-cases/unsafe-expect/unsafe-expect-1/vulnerable-example/src/lib.rs
@@ -68,4 +68,20 @@ mod tests {
         let balance = client.balance_of(&contract_id);
         assert_eq!(TOTAL_SUPPLY, balance.unwrap());
     }
+
+    #[test]
+    #[should_panic(expected = "could not get balance")]
+    fn balance_of_expect_works() {
+        // Given
+        let env = Env::default();
+        let contract_id = env.register_contract(None, UnsafeExpect);
+        let client = UnsafeExpectClient::new(&env, &contract_id);
+
+        // When - Balance not set
+
+        // Then
+        let _balance = client.balance_of(&contract_id);
+
+        // Test should panic
+    }
 }
diff --git a/test-cases/unsafe-expect/unsafe-expect-2/vulnerable-example/src/lib.rs b/test-cases/unsafe-expect/unsafe-expect-2/vulnerable-example/src/lib.rs
index 975d25d8..24e9d17f 100644
--- a/test-cases/unsafe-expect/unsafe-expect-2/vulnerable-example/src/lib.rs
+++ b/test-cases/unsafe-expect/unsafe-expect-2/vulnerable-example/src/lib.rs
@@ -66,4 +66,20 @@ mod tests {
         let balance = client.balance_of(&contract_id);
         assert_eq!(TOTAL_SUPPLY, balance.unwrap());
     }
+
+    #[test]
+    #[should_panic(expected = "could not get balance")]
+    fn balance_of_expect_works() {
+        // Given
+        let env = Env::default();
+        let contract_id = env.register_contract(None, UnsafeExpect);
+        let client = UnsafeExpectClient::new(&env, &contract_id);
+
+        // When - Balance not set
+
+        // Then
+        let _balance = client.balance_of(&contract_id);
+
+        // Test should panic
+    }
 }
diff --git a/test-cases/unsafe-expect/unsafe-expect-3/vulnerable-example/src/lib.rs b/test-cases/unsafe-expect/unsafe-expect-3/vulnerable-example/src/lib.rs
index 975d25d8..24e9d17f 100644
--- a/test-cases/unsafe-expect/unsafe-expect-3/vulnerable-example/src/lib.rs
+++ b/test-cases/unsafe-expect/unsafe-expect-3/vulnerable-example/src/lib.rs
@@ -66,4 +66,20 @@ mod tests {
         let balance = client.balance_of(&contract_id);
         assert_eq!(TOTAL_SUPPLY, balance.unwrap());
     }
+
+    #[test]
+    #[should_panic(expected = "could not get balance")]
+    fn balance_of_expect_works() {
+        // Given
+        let env = Env::default();
+        let contract_id = env.register_contract(None, UnsafeExpect);
+        let client = UnsafeExpectClient::new(&env, &contract_id);
+
+        // When - Balance not set
+
+        // Then
+        let _balance = client.balance_of(&contract_id);
+
+        // Test should panic
+    }
 }
diff --git a/test-cases/unsafe-expect/unsafe-expect-4/vulnerable-example/src/lib.rs b/test-cases/unsafe-expect/unsafe-expect-4/vulnerable-example/src/lib.rs
index 542ca1d4..29bbd9c3 100644
--- a/test-cases/unsafe-expect/unsafe-expect-4/vulnerable-example/src/lib.rs
+++ b/test-cases/unsafe-expect/unsafe-expect-4/vulnerable-example/src/lib.rs
@@ -74,4 +74,20 @@ mod tests {
         assert_eq!(TOTAL_SUPPLY, balances.0);
         assert_eq!(TOTAL_SUPPLY, balances.1);
     }
+
+    #[test]
+    #[should_panic(expected = "could not get balance")]
+    fn balance_of_expect_works() {
+        // Given
+        let env = Env::default();
+        let contract_id = env.register_contract(None, UnsafeExpect);
+        let client = UnsafeExpectClient::new(&env, &contract_id);
+
+        // When - Balance not set
+
+        // Then
+        let _balance = client.balance_of(&contract_id);
+
+        // Test should panic
+    }
 }
diff --git a/test-cases/unsafe-expect/unsafe-expect-5/vulnerable-example/src/lib.rs b/test-cases/unsafe-expect/unsafe-expect-5/vulnerable-example/src/lib.rs
index ae70acc0..fb8ebc5b 100644
--- a/test-cases/unsafe-expect/unsafe-expect-5/vulnerable-example/src/lib.rs
+++ b/test-cases/unsafe-expect/unsafe-expect-5/vulnerable-example/src/lib.rs
@@ -69,4 +69,20 @@ mod tests {
         // Then
         assert_eq!(TOTAL_SUPPLY + 2, balances.unwrap());
     }
+
+    #[test]
+    #[should_panic(expected = "could not get balance")]
+    fn balance_of_expect_works() {
+        // Given
+        let env = Env::default();
+        let contract_id = env.register_contract(None, UnsafeExpect);
+        let client = UnsafeExpectClient::new(&env, &contract_id);
+
+        // When - Balance not set
+
+        // Then
+        let _balance = client.balance_add(&contract_id);
+
+        // Test should panic
+    }
 }