From 174cd170c6456fc1b6250f1843631ba734679f50 Mon Sep 17 00:00:00 2001
From: Uxio Fuentefria <Uxio0@users.noreply.github.com>
Date: Fri, 13 Oct 2023 14:54:42 +0200
Subject: [PATCH] Add test for safe_creator

---
 safe_cli/safe_creator.py   |  5 +++--
 tests/test_entrypoint.py   | 14 ++++++--------
 tests/test_safe_creator.py | 38 ++++++++++++++++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/safe_cli/safe_creator.py b/safe_cli/safe_creator.py
index f50ac57e..3241097a 100644
--- a/safe_cli/safe_creator.py
+++ b/safe_cli/safe_creator.py
@@ -12,7 +12,7 @@
 from hexbytes import HexBytes
 from prompt_toolkit import print_formatted_text
 
-from gnosis.eth import EthereumClient
+from gnosis.eth import EthereumClient, EthereumTxSent
 from gnosis.eth.constants import NULL_ADDRESS
 from gnosis.eth.contracts import get_safe_V1_4_1_contract
 from gnosis.safe import ProxyFactory, Safe
@@ -104,7 +104,7 @@ def setup_argument_parser():
     return parser
 
 
-def main(*args, **kwargs):
+def main(*args, **kwargs) -> EthereumTxSent:
     parser = setup_argument_parser()
     print_formatted_text(pyfiglet.figlet_format("Safe Creator"))  # Print fancy text
     args = parser.parse_args()
@@ -216,3 +216,4 @@ def main(*args, **kwargs):
             f"will create safe={ethereum_tx_sent.contract_address}"
         )
         print_formatted_text(f"Tx parameters={ethereum_tx_sent.tx}")
+        return ethereum_tx_sent
diff --git a/tests/test_entrypoint.py b/tests/test_entrypoint.py
index 9ae6a761..c47cd516 100644
--- a/tests/test_entrypoint.py
+++ b/tests/test_entrypoint.py
@@ -19,15 +19,13 @@
 class SafeCliEntrypointTestCase(SafeCliTestCaseMixin, unittest.TestCase):
     random_safe_address = Account.create().address
 
-    @mock.patch(
-        "argparse.ArgumentParser.parse_args",
-        return_value=argparse.Namespace(
-            safe_address=random_safe_address,
-            node_url="http://localhost:8545",
-            history=True,
-        ),
-    )
+    @mock.patch("argparse.ArgumentParser.parse_args")
     def build_test_safe_cli(self, mock_parse_args: MagicMock):
+        mock_parse_args.return_value = argparse.Namespace(
+            safe_address=self.random_safe_address,
+            node_url=self.ethereum_node_url,
+            history=True,
+        )
         return build_safe_cli()
 
     @mock.patch.object(Safe, "retrieve_all_info")
diff --git a/tests/test_safe_creator.py b/tests/test_safe_creator.py
index ab04a215..670b089a 100644
--- a/tests/test_safe_creator.py
+++ b/tests/test_safe_creator.py
@@ -1,11 +1,45 @@
+import argparse
 import unittest
+from unittest import mock
+from unittest.mock import MagicMock
+
+from eth_account import Account
+
+from gnosis.safe import Safe
+
+from safe_cli.safe_creator import main
 
 from .safe_cli_test_case_mixin import SafeCliTestCaseMixin
 
 
 class TestSafeCreator(SafeCliTestCaseMixin, unittest.TestCase):
-    def test_main(self):
-        pass
+    @mock.patch(
+        "argparse.ArgumentParser.parse_args",
+    )
+    def test_main(self, mock_parse_args: MagicMock):
+        owner_account = self.get_ethereum_test_account()
+        owners = [Account.create().address]
+        threshold = 1
+        mock_parse_args.return_value = argparse.Namespace(
+            owners=owners,
+            threshold=threshold,
+            salt_nonce=4815,
+            node_url=self.ethereum_node_url,
+            private_key=owner_account.key.hex(),
+            safe_contract=self.safe_contract.address,
+            proxy_factory=self.proxy_factory_contract.address,
+            callback_handler=self.compatibility_fallback_handler.address,
+        )
+
+        safe_address = main().contract_address
+        safe = Safe(safe_address, self.ethereum_client)
+        safe_info = safe.retrieve_all_info()
+        self.assertEqual(safe_info.owners, owners)
+        self.assertEqual(safe_info.threshold, 1)
+        self.assertEqual(safe_info.master_copy, self.safe_contract.address)
+        self.assertEqual(
+            safe_info.fallback_handler, self.compatibility_fallback_handler.address
+        )
 
 
 if __name__ == "__main__":