From d8c479a2f3fc90733cda34e7e9ea760fee089c1f Mon Sep 17 00:00:00 2001
From: Andrei Marinica <andrei.marinica@multiversx.com>
Date: Fri, 21 Jun 2024 12:32:33 +0300
Subject: [PATCH] test address PartialEq implementation

---
 .../types/interaction/expr/test_address.rs    | 44 ++++++++++++++++++-
 .../types/interaction/expr/test_sc_address.rs | 39 +++++++++++++++-
 framework/scenario/tests/address_eq_test.rs   | 39 ++++++++++++++++
 3 files changed, 118 insertions(+), 4 deletions(-)
 create mode 100644 framework/scenario/tests/address_eq_test.rs

diff --git a/framework/base/src/types/interaction/expr/test_address.rs b/framework/base/src/types/interaction/expr/test_address.rs
index 88799c06d2..01d1dc819d 100644
--- a/framework/base/src/types/interaction/expr/test_address.rs
+++ b/framework/base/src/types/interaction/expr/test_address.rs
@@ -6,11 +6,13 @@ use crate::{
     abi::TypeAbiFrom,
     api::ManagedTypeApi,
     types::{
-        AnnotatedValue, ManagedAddress, ManagedBuffer, TxEnv, TxFrom, TxFromSpecified, TxTo,
-        TxToSpecified,
+        heap::Address, AnnotatedValue, ManagedAddress, ManagedBuffer, TxEnv, TxFrom,
+        TxFromSpecified, TxTo, TxToSpecified,
     },
 };
 
+use super::TestSCAddress;
+
 const ADDRESS_PREFIX: &str = "address:";
 
 /// Encodes a dummy address, to be used for tests.
@@ -40,12 +42,50 @@ impl<'a> TestAddress<'a> {
         result
     }
 
+    pub fn to_address(&self) -> Address {
+        self.eval_to_array().into()
+    }
+
+    pub fn to_managed_address<Api: ManagedTypeApi>(&self) -> ManagedAddress<Api> {
+        self.eval_to_array().into()
+    }
+
     #[cfg(feature = "alloc")]
     pub fn eval_to_expr(&self) -> alloc::string::String {
         alloc::format!("{ADDRESS_PREFIX}{}", self.name)
     }
 }
 
+impl<'a, 'b> PartialEq<TestSCAddress<'b>> for TestAddress<'a> {
+    fn eq(&self, other: &TestSCAddress) -> bool {
+        self.to_address() == other.to_address()
+    }
+}
+
+impl<'a> PartialEq<Address> for TestAddress<'a> {
+    fn eq(&self, other: &Address) -> bool {
+        &self.to_address() == other
+    }
+}
+
+impl<'a> PartialEq<TestAddress<'a>> for Address {
+    fn eq(&self, other: &TestAddress<'a>) -> bool {
+        self == &other.to_address()
+    }
+}
+
+impl<'a, Api: ManagedTypeApi> PartialEq<ManagedAddress<Api>> for TestAddress<'a> {
+    fn eq(&self, other: &ManagedAddress<Api>) -> bool {
+        self.to_address() == other.to_address()
+    }
+}
+
+impl<'a, Api: ManagedTypeApi> PartialEq<TestAddress<'a>> for ManagedAddress<Api> {
+    fn eq(&self, other: &TestAddress<'a>) -> bool {
+        self.to_address() == other.to_address()
+    }
+}
+
 impl<'a, Env> AnnotatedValue<Env, ManagedAddress<Env::Api>> for TestAddress<'a>
 where
     Env: TxEnv,
diff --git a/framework/base/src/types/interaction/expr/test_sc_address.rs b/framework/base/src/types/interaction/expr/test_sc_address.rs
index bf01de38bf..564157fb94 100644
--- a/framework/base/src/types/interaction/expr/test_sc_address.rs
+++ b/framework/base/src/types/interaction/expr/test_sc_address.rs
@@ -11,6 +11,8 @@ use crate::{
     },
 };
 
+use super::TestAddress;
+
 const SC_PREFIX: &str = "sc:";
 const VM_TYPE_LEN: usize = 2;
 const DEFAULT_VM_TYPE: &[u8] = &[5, 0];
@@ -48,8 +50,41 @@ where
 
 impl<'a> TestSCAddress<'a> {
     pub fn to_address(&self) -> Address {
-        let expr: [u8; 32] = self.eval_to_array();
-        expr.into()
+        self.eval_to_array().into()
+    }
+
+    pub fn to_managed_address<Api: ManagedTypeApi>(&self) -> ManagedAddress<Api> {
+        self.eval_to_array().into()
+    }
+}
+
+impl<'a, 'b> PartialEq<TestAddress<'b>> for TestSCAddress<'a> {
+    fn eq(&self, other: &TestAddress) -> bool {
+        self.to_address() == other.to_address()
+    }
+}
+
+impl<'a> PartialEq<Address> for TestSCAddress<'a> {
+    fn eq(&self, other: &Address) -> bool {
+        &self.to_address() == other
+    }
+}
+
+impl<'a> PartialEq<TestSCAddress<'a>> for Address {
+    fn eq(&self, other: &TestSCAddress<'a>) -> bool {
+        self == &other.to_address()
+    }
+}
+
+impl<'a, Api: ManagedTypeApi> PartialEq<ManagedAddress<Api>> for TestSCAddress<'a> {
+    fn eq(&self, other: &ManagedAddress<Api>) -> bool {
+        self.to_address() == other.to_address()
+    }
+}
+
+impl<'a, Api: ManagedTypeApi> PartialEq<TestSCAddress<'a>> for ManagedAddress<Api> {
+    fn eq(&self, other: &TestSCAddress<'a>) -> bool {
+        self.to_address() == other.to_address()
     }
 }
 
diff --git a/framework/scenario/tests/address_eq_test.rs b/framework/scenario/tests/address_eq_test.rs
new file mode 100644
index 0000000000..af78e6778b
--- /dev/null
+++ b/framework/scenario/tests/address_eq_test.rs
@@ -0,0 +1,39 @@
+use multiversx_sc::types::TestAddress;
+use multiversx_sc_scenario::api::StaticApi;
+
+const ALICE: TestAddress = TestAddress::new("alice");
+const SC_ADDR: TestAddress = TestAddress::new("sc");
+
+#[test]
+fn test_address_eq() {
+    let alice2 = TestAddress::new("alice");
+
+    assert_eq!(ALICE, alice2);
+    assert_eq!(ALICE, alice2.to_address());
+    assert_eq!(ALICE.to_address(), alice2);
+    assert_eq!(ALICE.to_address(), alice2.to_address());
+
+    assert_eq!(ALICE, alice2.to_managed_address::<StaticApi>());
+    assert_eq!(ALICE.to_managed_address::<StaticApi>(), alice2);
+    assert_eq!(
+        ALICE.to_managed_address::<StaticApi>(),
+        alice2.to_managed_address::<StaticApi>()
+    );
+}
+
+#[test]
+fn test_sc_address_eq() {
+    let sc2 = TestAddress::new("sc");
+
+    assert_eq!(SC_ADDR, sc2);
+    assert_eq!(SC_ADDR, sc2.to_address());
+    assert_eq!(SC_ADDR.to_address(), sc2);
+    assert_eq!(SC_ADDR.to_address(), sc2.to_address());
+
+    assert_eq!(SC_ADDR, sc2.to_managed_address::<StaticApi>());
+    assert_eq!(SC_ADDR.to_managed_address::<StaticApi>(), sc2);
+    assert_eq!(
+        SC_ADDR.to_managed_address::<StaticApi>(),
+        sc2.to_managed_address::<StaticApi>()
+    );
+}