Add cmdline/CLI tests (#258)

pyicloud/cmdline.py

@@ -8,12 +8,18 @@ from __future__ import print_function
 import argparse
 import pickle
 import sys
+import six
 
-from click import confirm, prompt
+from click import confirm
 
-import pyicloud
+from pyicloud import PyiCloudService
+from pyicloud.exceptions import PyiCloudFailedLoginException
 from . import utils
 
+if six.PY2:
+    input = raw_input  # pylint: disable=redefined-builtin,invalid-name,undefined-variable
+else:
+    input = input   # pylint: disable=bad-option-value,self-assigning-variable,invalid-name
 
 DEVICE_ERROR = (
     "Please use the --device switch to indicate which device to use."
@@ -21,19 +27,20 @@ DEVICE_ERROR = (
 
 
 def create_pickled_data(idevice, filename):
-    """This helper will output the idevice to a pickled file named
+    """
+    This helper will output the idevice to a pickled file named
     after the passed filename.
 
     This allows the data to be used without resorting to screen / pipe
-    scrapping."""
+    scrapping.
-    location = filename
+    """
-    pickle_file = open(location, 'wb')
+    pickle_file = open(filename, 'wb')
     pickle.dump(idevice.content, pickle_file, protocol=pickle.HIGHEST_PROTOCOL)
     pickle_file.close()
 
 
 def main(args=None):
-    """Main commandline entrypoint"""
+    """Main commandline entrypoint."""
     if args is None:
         args = sys.argv[1:]
 
@@ -94,7 +101,7 @@ def main(args=None):
         help="Retrieve Location for the iDevice (non-exclusive).",
     )
 
-    #   Restrict actions to a specific devices UID / DID
+    # Restrict actions to a specific devices UID / DID
     parser.add_argument(
         "--device",
         action="store",
@@ -103,7 +110,7 @@ def main(args=None):
         help="Only effect this device",
     )
 
-    #   Trigger Sound Alert
+    # Trigger Sound Alert
     parser.add_argument(
         "--sound",
         action="store_true",
@@ -112,7 +119,7 @@ def main(args=None):
         help="Play a sound on the device",
     )
 
-    #   Trigger Message w/Sound Alert
+    # Trigger Message w/Sound Alert
     parser.add_argument(
         "--message",
         action="store",
@@ -121,7 +128,7 @@ def main(args=None):
         help="Optional Text Message to display with a sound",
     )
 
-    #   Trigger Message (without Sound) Alert
+    # Trigger Message (without Sound) Alert
     parser.add_argument(
         "--silentmessage",
         action="store",
@@ -130,7 +137,7 @@ def main(args=None):
         help="Optional Text Message to display with no sounds",
     )
 
-    #   Lost Mode
+    # Lost Mode
     parser.add_argument(
         "--lostmode",
         action="store_true",
@@ -160,7 +167,7 @@ def main(args=None):
         help="Forcibly display this message when activating lost mode.",
     )
 
-    #   Output device data to an pickle file
+    # Output device data to an pickle file
     parser.add_argument(
         "--outputfile",
         action="store_true",
@@ -194,42 +201,44 @@ def main(args=None):
             parser.error('No password supplied')
 
         try:
-            api = pyicloud.PyiCloudService(
+            api = PyiCloudService(
                 username.strip(),
                 password.strip()
             )
             if (
                 not utils.password_exists_in_keyring(username) and
                 command_line.interactive and
-                confirm("Save password in keyring? ")
+                confirm("Save password in keyring?")
             ):
                 utils.store_password_in_keyring(username, password)
 
             if api.requires_2sa:
-                print("Two-step authentication required.",
+                print("\nTwo-step authentication required.",
-                      "Your trusted devices are:")
+                      "\nYour trusted devices are:")
 
                 devices = api.trusted_devices
                 for i, device in enumerate(devices):
-                    print("  %s: %s" % (
+                    print("    %s: %s" % (
                         i, device.get(
                             'deviceName',
                             "SMS to %s" % device.get('phoneNumber'))))
 
-                device = prompt('Which device would you like to use?',
+                print('\nWhich device would you like to use?')
-                                      default=0)
+                device = int(input('(number) --> '))
                 device = devices[device]
                 if not api.send_verification_code(device):
                     print("Failed to send verification code")
                     sys.exit(1)
 
-                code = prompt('Please enter validation code')
+                print('\nPlease enter validation code')
+                code = input('(string) --> ')
                 if not api.validate_verification_code(device, code):
                     print("Failed to verify verification code")
                     sys.exit(1)
 
+                print('')
             break
-        except pyicloud.exceptions.PyiCloudFailedLoginException:
+        except PyiCloudFailedLoginException:
             # If they have a stored password; we just used it and
             # it did not work; let's delete it if there is one.
             if utils.password_exists_in_keyring(username):
@@ -282,7 +291,7 @@ def main(args=None):
                 print("Device Class  - %s" % contents["deviceClass"])
                 print("Device Model  - %s" % contents["deviceModel"])
 
-            #   Play a Sound on a device
+            # Play a Sound on a device
             if command_line.sound:
                 if command_line.device_id:
                     dev.play_sound()
@@ -294,7 +303,7 @@ def main(args=None):
                         )
                     )
 
-            #   Display a Message on the device
+            # Display a Message on the device
             if command_line.message:
                 if command_line.device_id:
                     dev.display_message(
@@ -311,7 +320,7 @@ def main(args=None):
                         )
                     )
 
-            #   Display a Silent Message on the device
+            # Display a Silent Message on the device
             if command_line.silentmessage:
                 if command_line.device_id:
                     dev.display_message(
@@ -328,7 +337,7 @@ def main(args=None):
                         )
                     )
 
-            #   Enable Lost mode
+            # Enable Lost mode
             if command_line.lostmode:
                 if command_line.device_id:
                     dev.lost_device(
@@ -344,6 +353,7 @@ def main(args=None):
                             DEVICE_ERROR
                         )
                     )
+    sys.exit(0)
 
 if __name__ == '__main__':
     main()

tests/__init__.py

@@ -1 +1,190 @@
 """Library tests."""
+
+from pyicloud import base
+from pyicloud.exceptions import PyiCloudFailedLoginException
+from pyicloud.services.findmyiphone import FindMyiPhoneServiceManager, AppleDevice
+
+
+AUTHENTICATED_USER = "authenticated_user"
+REQUIRES_2SA_USER = "requires_2sa_user"
+VALID_USERS = [AUTHENTICATED_USER, REQUIRES_2SA_USER]
+
+
+class PyiCloudServiceMock(base.PyiCloudService):
+    """Mocked PyiCloudService."""
+    def __init__(
+            self,
+            apple_id,
+            password=None,
+            cookie_directory=None,
+            verify=True,
+            client_id=None,
+            with_family=True
+        ):
+        base.PyiCloudService.__init__(self, apple_id, password, cookie_directory, verify, client_id, with_family)
+        base.FindMyiPhoneServiceManager = FindMyiPhoneServiceManagerMock
+
+    def authenticate(self):
+        if not self.user.get("apple_id") or self.user.get("apple_id") not in VALID_USERS:
+            raise PyiCloudFailedLoginException("Invalid email/password combination.", None)
+        if not self.user.get("password") or self.user.get("password") != "valid_pass":
+            raise PyiCloudFailedLoginException("Invalid email/password combination.", None)
+        
+        self.params.update({'dsid': 'ID'})
+        self._webservices = {
+            'account': {
+                'url': 'account_url',
+            },
+            'findme': {
+                'url': 'findme_url',
+            },
+            'calendar': {
+                'url': 'calendar_url',
+            },
+            'contacts': {
+                'url': 'contacts_url',
+            },
+            'reminders': {
+                'url': 'reminders_url',
+            }
+        }
+
+    @property
+    def requires_2sa(self):
+        return self.user["apple_id"] is REQUIRES_2SA_USER
+
+    @property
+    def trusted_devices(self):
+        return  [
+            {"deviceType": "SMS", "areaCode": "", "phoneNumber": "*******58", "deviceId": "1"}
+        ]
+    
+    def send_verification_code(self, device):
+        return device
+    
+    def validate_verification_code(self, device, code):
+        if not device or code != 0:
+            self.user["apple_id"] = AUTHENTICATED_USER
+        self.authenticate()
+        return not self.requires_2sa
+
+
+IPHONE_DEVICE_ID = "X1x/X&x="
+IPHONE_DEVICE = AppleDevice(
+    {
+        "msg": {
+            "strobe": False,
+            "userText": False,
+            "playSound": True,
+            "vibrate": True,
+            "createTimestamp": 1568031021347,
+            "statusCode": "200"
+        },
+        "canWipeAfterLock": True,
+        "baUUID": "",
+        "wipeInProgress": False,
+        "lostModeEnabled": False,
+        "activationLocked": True,
+        "passcodeLength": 6,
+        "deviceStatus": "200",
+        "deviceColor": "1-6-0",
+        "features": {
+            "MSG": True,
+            "LOC": True,
+            "LLC": False,
+            "CLK": False,
+            "TEU": True,
+            "LMG": False,
+            "SND": True,
+            "CLT": False,
+            "LKL": False,
+            "SVP": False,
+            "LST": True,
+            "LKM": False,
+            "WMG": True,
+            "SPN": False,
+            "XRM": False,
+            "PIN": False,
+            "LCK": True,
+            "REM": False,
+            "MCS": False,
+            "CWP": False,
+            "KEY": False,
+            "KPD": False,
+            "WIP": True
+        },
+        "lowPowerMode": True,
+        "rawDeviceModel": "iPhone11,8",
+        "id": IPHONE_DEVICE_ID,
+        "remoteLock": None,
+        "isLocating": True,
+        "modelDisplayName": "iPhone",
+        "lostTimestamp": "",
+        "batteryLevel": 0.47999998927116394,
+        "mesg": None,
+        "locationEnabled": True,
+        "lockedTimestamp": None,
+        "locFoundEnabled": False,
+        "snd": {
+            "createTimestamp": 1568031021347,
+            "statusCode": "200"
+        },
+        "fmlyShare": False,
+        "lostDevice": {
+            "stopLostMode": False,
+            "emailUpdates": False,
+            "userText": True,
+            "sound": False,
+            "ownerNbr": "",
+            "text": "",
+            "createTimestamp": 1558383841233,
+            "statusCode": "2204"
+        },
+        "lostModeCapable": True,
+        "wipedTimestamp": None,
+        "deviceDisplayName": "iPhone XR",
+        "prsId": None,
+        "audioChannels": [],
+        "locationCapable": True,
+        "batteryStatus": "NotCharging",
+        "trackingInfo": None,
+        "name": "Quentin's iPhone",
+        "isMac": False,
+        "thisDevice": False,
+        "deviceClass": "iPhone",
+        "location": {
+            "isOld": False,
+            "isInaccurate": False,
+            "altitude": 0.0,
+            "positionType": "GPS",
+            "latitude": 46.012345678,
+            "floorLevel": 0,
+            "horizontalAccuracy": 12.012345678,
+            "locationType": "",
+            "timeStamp": 1568827039692,
+            "locationFinished": False,
+            "verticalAccuracy": 0.0,
+            "longitude": 5.012345678
+        },
+        "deviceModel": "iphoneXR-1-6-0",
+        "maxMsgChar": 160,
+        "darkWake": False,
+        "remoteWipe": None
+    },
+    None,
+    None,
+    None
+)
+
+DEVICES = {
+    IPHONE_DEVICE_ID: IPHONE_DEVICE,
+}
+
+
+class FindMyiPhoneServiceManagerMock(FindMyiPhoneServiceManager):
+    """Mocked FindMyiPhoneServiceManager."""
+    def __init__(self, service_root, session, params, with_family=False):
+        FindMyiPhoneServiceManager.__init__(self, service_root, session, params, with_family)
+
+    def refresh_client(self):
+        self._devices = DEVICES

tests/test_cmdline.py

@@ -0,0 +1,101 @@
+"""Cmdline tests."""
+from pyicloud import cmdline
+from . import PyiCloudServiceMock, AUTHENTICATED_USER, REQUIRES_2SA_USER, DEVICES
+
+import os
+import sys
+import pickle
+import pytest
+from unittest import TestCase
+if sys.version_info >= (3, 3):
+    from unittest.mock import patch  # pylint: disable=no-name-in-module,import-error
+else:
+    from mock import patch
+
+class TestCmdline(TestCase):
+    """Cmdline test cases."""
+    main = None
+
+    def setUp(self):
+        cmdline.PyiCloudService = PyiCloudServiceMock
+        self.main = cmdline.main
+
+    def test_no_arg(self):
+        """Test no args."""
+        with pytest.raises(SystemExit, match="2"):
+            self.main()
+
+        with pytest.raises(SystemExit, match="2"):
+            self.main(None)
+
+        with pytest.raises(SystemExit, match="2"):
+            self.main([])
+
+    def test_help(self):
+        """Test the help command."""
+        with pytest.raises(SystemExit, match="0"):
+            self.main(['--help'])
+
+    def test_username(self):
+        """Test the username command."""
+        # No username supplied
+        with pytest.raises(SystemExit, match="2"):
+            self.main(['--username'])
+
+    @patch("getpass.getpass")
+    def test_username_password_invalid(self, mock_getpass):
+        """Test username and password commands."""
+        # No password supplied
+        mock_getpass.return_value = None
+        with pytest.raises(SystemExit, match="2"):
+            self.main(['--username', 'invalid_user'])
+
+        # Bad username or password
+        mock_getpass.return_value = "invalid_pass"
+        with pytest.raises(RuntimeError, match="Bad username or password for invalid_user"):
+            self.main(['--username', 'invalid_user'])
+
+        # We should not use getpass for this one, but we reset the password at login fail
+        with pytest.raises(RuntimeError, match="Bad username or password for invalid_user"):
+            self.main(['--username', 'invalid_user', '--password', 'invalid_pass'])
+        
+
+    @patch('pyicloud.cmdline.input')
+    def test_username_password_requires_2sa(self, mock_input):
+        """Test username and password commands."""
+        # Valid connection for the first time
+        mock_input.return_value = "0"
+        with pytest.raises(SystemExit, match="0"):
+            self.main([
+                '--username', REQUIRES_2SA_USER,
+                '--password', 'valid_pass',
+                '--non-interactive',
+            ])
+
+    def test_device_outputfile(self):
+        """Test the outputfile command."""
+        with pytest.raises(SystemExit, match="0"):
+            self.main([
+                '--username', AUTHENTICATED_USER,
+                '--password', 'valid_pass',
+                '--non-interactive',
+                '--outputfile'
+            ])
+
+        for key in DEVICES:
+            file_name = DEVICES[key].content['name'].strip().lower() + ".fmip_snapshot"
+
+            pickle_file = open(file_name, "rb")
+            assert pickle_file
+
+            contents = []
+            with pickle_file as opened_file:
+                while True:
+                    try:
+                        contents.append(pickle.load(opened_file))
+                    except EOFError:
+                        break
+            assert contents == [DEVICES[key].content]
+
+            pickle_file.close()
+            os.remove(file_name)

tests/test_sanity.py

@@ -1,12 +0,0 @@
-"""Sanity test."""
-from unittest2 import TestCase
-
-from pyicloud.cmdline import main
-
-
-class SanityTestCase(TestCase):
-    """Sanity test."""
-    def test_basic_sanity(self):
-        """Sanity test."""
-        with self.assertRaises(SystemExit):
-            main(['--help'])