diff --git a/ape_safe/_cli/click_ext.py b/ape_safe/_cli/click_ext.py
index 271bb68..672bcd3 100644
--- a/ape_safe/_cli/click_ext.py
+++ b/ape_safe/_cli/click_ext.py
@@ -50,6 +50,7 @@ def _safe_callback(ctx, param, value):
 
 
 safe_option = click.option("--safe", callback=_safe_callback)
+safe_argument = click.argument("safe", callback=_safe_callback)
 
 
 def _txn_ids_callback(ctx, param, value):
diff --git a/ape_safe/_cli/safe_mgmt.py b/ape_safe/_cli/safe_mgmt.py
index 216551f..71442b2 100644
--- a/ape_safe/_cli/safe_mgmt.py
+++ b/ape_safe/_cli/safe_mgmt.py
@@ -4,7 +4,7 @@
 from ape.types import AddressType
 from eth_typing import ChecksumAddress
 
-from ape_safe._cli.click_ext import SafeCliContext, safe_cli_ctx, safe_option
+from ape_safe._cli.click_ext import SafeCliContext, safe_argument, safe_cli_ctx
 from ape_safe.client import ExecutedTxData
 
 
@@ -84,7 +84,7 @@ def add(cli_ctx: SafeCliContext, ecosystem, network, address, alias):
 
 @click.command()
 @safe_cli_ctx()
-@safe_option
+@safe_argument
 def remove(cli_ctx: SafeCliContext, safe):
     """
     Stop tracking a locally-tracked Safe
diff --git a/tests/test_account.py b/tests/functional/test_account.py
similarity index 100%
rename from tests/test_account.py
rename to tests/functional/test_account.py
diff --git a/tests/test_multisend.py b/tests/functional/test_multisend.py
similarity index 100%
rename from tests/test_multisend.py
rename to tests/functional/test_multisend.py
diff --git a/tests/test_utils.py b/tests/functional/test_utils.py
similarity index 100%
rename from tests/test_utils.py
rename to tests/functional/test_utils.py
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
new file mode 100644
index 0000000..798ca67
--- /dev/null
+++ b/tests/integration/conftest.py
@@ -0,0 +1,28 @@
+import shutil
+
+import pytest
+from click.testing import CliRunner
+
+from ape_safe._cli import cli as CLI
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+@pytest.fixture
+def cli():
+    return CLI
+
+
+@pytest.fixture
+def no_safes(data_folder):
+    shutil.rmtree(data_folder, ignore_errors=True)
+
+
+@pytest.fixture
+def one_safe(data_folder, safes, safe):
+    shutil.rmtree(data_folder)
+    safes.save_account(safe.alias, safe.address)
+    return safes.load_account(safe.alias)
diff --git a/tests/integration/test_pending_cli.py b/tests/integration/test_pending_cli.py
new file mode 100644
index 0000000..40792fe
--- /dev/null
+++ b/tests/integration/test_pending_cli.py
@@ -0,0 +1,3 @@
+def test_help(runner, cli):
+    result = runner.invoke(cli, ["--help"], catch_exceptions=False)
+    assert result.exit_code == 0, result.output
diff --git a/tests/integration/test_safe_mgmt_cli.py b/tests/integration/test_safe_mgmt_cli.py
new file mode 100644
index 0000000..8532b55
--- /dev/null
+++ b/tests/integration/test_safe_mgmt_cli.py
@@ -0,0 +1,25 @@
+def test_list_no_safes(runner, cli, no_safes):
+    result = runner.invoke(cli, ["list"], catch_exceptions=False)
+    assert result.exit_code == 0, result.output
+    assert "No Safes found" in result.output
+
+
+def test_list_one_safe(runner, cli, one_safe):
+    result = runner.invoke(cli, ["list"], catch_exceptions=False)
+    assert result.exit_code == 0, result.output
+    assert "Found 1 Safe" in result.output
+    assert "0x5FbDB2315678afecb367f032d93F642f64180aa3" in result.output
+
+
+def test_add_safe(runner, cli, no_safes, safe):
+    result = runner.invoke(
+        cli, ["add", safe.address, safe.alias], catch_exceptions=False, input="y\n"
+    )
+    assert result.exit_code == 0, result.output
+    assert "SUCCESS" in result.output, result.output
+
+
+def test_remove_safe(runner, cli, one_safe, safe):
+    result = runner.invoke(cli, ["remove", safe.alias], catch_exceptions=False, input="y\n")
+    assert result.exit_code == 0, result.output
+    assert "SUCCESS" in result.output, result.output
diff --git a/tests/test_cli.py b/tests/test_cli.py
deleted file mode 100644
index c702908..0000000
--- a/tests/test_cli.py
+++ /dev/null
@@ -1,41 +0,0 @@
-import shutil
-
-import pytest
-from click.testing import CliRunner
-
-from ape_safe._cli import cli as CLI
-
-
-@pytest.fixture
-def runner():
-    return CliRunner()
-
-
-@pytest.fixture
-def cli():
-    return CLI
-
-
-@pytest.fixture
-def no_safes(data_folder):
-    shutil.rmtree(data_folder, ignore_errors=True)
-
-
-@pytest.fixture
-def one_safe(data_folder, safes, safe):
-    shutil.rmtree(data_folder)
-    safes.save_account(safe.alias, safe.address)
-    return safes.load_account(safe.alias)
-
-
-def test_list_no_safes(runner, cli, no_safes):
-    result = runner.invoke(cli, ["list"], catch_exceptions=False)
-    assert result.exit_code == 0, result.output
-    assert "No Safes found" in result.output
-
-
-def test_list_one_safe(runner, cli, one_safe):
-    result = runner.invoke(cli, ["list"], catch_exceptions=False)
-    assert result.exit_code == 0, result.output
-    assert "Found 1 Safe" in result.output
-    assert "0x5FbDB2315678afecb367f032d93F642f64180aa3" in result.output