From 59e14d37411b977a8fb9811fdc9678dfeed4aa9c Mon Sep 17 00:00:00 2001
From: Mike Kinney <mike.kinney@gmail.com>
Date: Sun, 12 Dec 2021 14:37:41 -0800
Subject: [PATCH] add main unit tests for --set-team

---
 meshtastic/__main__.py        | 10 +++---
 meshtastic/tests/test_main.py | 67 +++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/meshtastic/__main__.py b/meshtastic/__main__.py
index 19a6393..5b7a110 100644
--- a/meshtastic/__main__.py
+++ b/meshtastic/__main__.py
@@ -16,7 +16,7 @@ import meshtastic.test
 from .tcp_interface import TCPInterface
 from .ble_interface import BLEInterface
 from . import remote_hardware
-from . import portnums_pb2, channel_pb2, mesh_pb2, radioconfig_pb2
+from . import portnums_pb2, channel_pb2, radioconfig_pb2
 from .globals import Globals
 
 """We only import the tunnel code if we are on a platform that can run it"""
@@ -219,14 +219,14 @@ def onConnected(interface):
         if args.set_team:
             closeNow = True
             try:
-                v_team = mesh_pb2.Team.Value(args.set_team.upper())
+                v_team = meshtastic.mesh_pb2.Team.Value(args.set_team.upper())
             except ValueError:
                 v_team = 0
                 print(f"ERROR: Team \'{args.set_team}\' not found.")
-                print("Try a team name from the list below, or CLEAR for unaffiliated:")
-                print(mesh_pb2.Team.keys())
+                print("Try a team name from the sorted list below, or use 'CLEAR' for unaffiliated:")
+                print(sorted(meshtastic.mesh_pb2.Team.keys()))
             else:
-                print(f"Setting team to {mesh_pb2.Team.Name(v_team)}")
+                print(f"Setting team to {meshtastic.mesh_pb2.Team.Name(v_team)}")
                 getNode().setOwner(team=v_team)
 
         if args.set_ham:
diff --git a/meshtastic/tests/test_main.py b/meshtastic/tests/test_main.py
index 7f7bb7b..d9f0006 100644
--- a/meshtastic/tests/test_main.py
+++ b/meshtastic/tests/test_main.py
@@ -537,3 +537,70 @@ def test_main_setalt(capsys):
         # TODO: Why does this not work? assert re.search(r'inside mocked writeConfig', out, re.MULTILINE)
         assert err == ''
         mo.assert_called()
+
+
+@pytest.mark.unit
+def test_main_set_team_valid(capsys):
+    """Test --set-team"""
+    sys.argv = ['', '--set-team', 'CYAN']
+    args = sys.argv
+    parser = None
+    parser = argparse.ArgumentParser()
+    our_globals = Globals.getInstance()
+    our_globals.set_parser(parser)
+    our_globals.set_args(args)
+    our_globals.set_target_node(None)
+
+    mocked_node = MagicMock(autospec=Node)
+    def mock_setOwner(team):
+        print('inside mocked setOwner')
+    mocked_node.setOwner.side_effect = mock_setOwner
+
+    iface = MagicMock(autospec=SerialInterface)
+    iface.localNode.return_value = mocked_node
+
+    with patch('meshtastic.serial_interface.SerialInterface', return_value=iface) as mo:
+        with patch('meshtastic.mesh_pb2.Team') as mm:
+            mm.Name.return_value = 'FAKENAME'
+            mm.Value.return_value = 'FAKEVAL'
+            main()
+            out, err = capsys.readouterr()
+            print('out:', out)
+            print('err:', err)
+            assert re.search(r'Connected to radio', out, re.MULTILINE)
+            assert re.search(r'Setting team to', out, re.MULTILINE)
+            assert err == ''
+            mo.assert_called()
+            mm.Name.assert_called()
+            mm.Value.assert_called()
+
+
+@pytest.mark.unit
+def test_main_set_team_invalid(capsys):
+    """Test --set-team using an invalid team name"""
+    sys.argv = ['', '--set-team', 'NOTCYAN']
+    args = sys.argv
+    parser = None
+    parser = argparse.ArgumentParser()
+    our_globals = Globals.getInstance()
+    our_globals.set_parser(parser)
+    our_globals.set_args(args)
+    our_globals.set_target_node(None)
+
+    iface = MagicMock(autospec=SerialInterface)
+
+    def throw_an_exception(exc):
+        raise ValueError("Fake exception.")
+
+    with patch('meshtastic.serial_interface.SerialInterface', return_value=iface) as mo:
+        with patch('meshtastic.mesh_pb2.Team') as mm:
+            mm.Value.side_effect = throw_an_exception
+            main()
+            out, err = capsys.readouterr()
+            print('out:', out)
+            print('err:', err)
+            assert re.search(r'Connected to radio', out, re.MULTILINE)
+            assert re.search(r'ERROR: Team', out, re.MULTILINE)
+            assert err == ''
+            mo.assert_called()
+            mm.Value.assert_called()