pyicloud/tests/__init__.py

"""Library tests."""
import json
from requests import Response

from pyicloud import base

from .const import (
    AUTHENTICATED_USER,
    REQUIRES_2FA_USER,
    REQUIRES_2FA_TOKEN,
    VALID_TOKEN,
    VALID_USERS,
    VALID_PASSWORD,
    VALID_COOKIE,
    VALID_2FA_CODE,
    VALID_TOKENS,
)
from .const_login import (
    AUTH_OK,
    LOGIN_WORKING,
    LOGIN_2FA,
    TRUSTED_DEVICES,
    TRUSTED_DEVICE_1,
    VERIFICATION_CODE_OK,
    VERIFICATION_CODE_KO,
)
from .const_account import ACCOUNT_DEVICES_WORKING, ACCOUNT_STORAGE_WORKING
from .const_account_family import ACCOUNT_FAMILY_WORKING
from .const_drive import (
    DRIVE_FOLDER_WORKING,
    DRIVE_ROOT_INVALID,
    DRIVE_SUBFOLDER_WORKING,
    DRIVE_ROOT_WORKING,
    DRIVE_FILE_DOWNLOAD_WORKING,
)
from .const_findmyiphone import FMI_FAMILY_WORKING


class ResponseMock(Response):
    """Mocked Response."""

    def __init__(self, result, status_code=200, **kwargs):
        """Set up response mock."""
        Response.__init__(self)
        self.result = result
        self.status_code = status_code
        self.raw = kwargs.get("raw")
        self.headers = kwargs.get("headers", {})

    @property
    def text(self):
        """Return text."""
        return json.dumps(self.result)


class PyiCloudSessionMock(base.PyiCloudSession):
    """Mocked PyiCloudSession."""

    def request(self, method, url, **kwargs):
        """Make the request."""
        params = kwargs.get("params")
        headers = kwargs.get("headers")
        data = json.loads(kwargs.get("data", "{}"))

        # Login
        if self.service.SETUP_ENDPOINT in url:
            if "accountLogin" in url and method == "POST":
                if data.get("dsWebAuthToken") not in VALID_TOKENS:
                    self._raise_error(None, "Unknown reason")
                if data.get("dsWebAuthToken") == REQUIRES_2FA_TOKEN:
                    return ResponseMock(LOGIN_2FA)
                return ResponseMock(LOGIN_WORKING)

            if "listDevices" in url and method == "GET":
                return ResponseMock(TRUSTED_DEVICES)

            if "sendVerificationCode" in url and method == "POST":
                if data == TRUSTED_DEVICE_1:
                    return ResponseMock(VERIFICATION_CODE_OK)
                return ResponseMock(VERIFICATION_CODE_KO)

            if "validateVerificationCode" in url and method == "POST":
                TRUSTED_DEVICE_1.update({"verificationCode": "0", "trustBrowser": True})
                if data == TRUSTED_DEVICE_1:
                    self.service.user["apple_id"] = AUTHENTICATED_USER
                    return ResponseMock(VERIFICATION_CODE_OK)
                self._raise_error(None, "FOUND_CODE")

            if "validate" in url and method == "POST":
                if headers.get("X-APPLE-WEBAUTH-TOKEN") == VALID_COOKIE:
                    return ResponseMock(LOGIN_WORKING)
                self._raise_error(None, "Session expired")

        if self.service.AUTH_ENDPOINT in url:
            if "signin" in url and method == "POST":
                if (
                    data.get("accountName") not in VALID_USERS
                    or data.get("password") != VALID_PASSWORD
                ):
                    self._raise_error(None, "Unknown reason")
                if data.get("accountName") == REQUIRES_2FA_USER:
                    self.service.session_data["session_token"] = REQUIRES_2FA_TOKEN
                    return ResponseMock(AUTH_OK)

                self.service.session_data["session_token"] = VALID_TOKEN
                return ResponseMock(AUTH_OK)

            if "securitycode" in url and method == "POST":
                if data.get("securityCode", {}).get("code") != VALID_2FA_CODE:
                    self._raise_error(None, "Incorrect code")

                self.service.session_data["session_token"] = VALID_TOKEN
                return ResponseMock("", status_code=204)

            if "trust" in url and method == "GET":
                return ResponseMock("", status_code=204)

        # Account
        if "device/getDevices" in url and method == "GET":
            return ResponseMock(ACCOUNT_DEVICES_WORKING)
        if "family/getFamilyDetails" in url and method == "GET":
            return ResponseMock(ACCOUNT_FAMILY_WORKING)
        if "setup/ws/1/storageUsageInfo" in url and method == "GET":
            return ResponseMock(ACCOUNT_STORAGE_WORKING)

        # Drive
        if (
            "retrieveItemDetailsInFolders" in url
            and method == "POST"
            and data[0].get("drivewsid")
        ):
            if data[0].get("drivewsid") == "FOLDER::com.apple.CloudDocs::root":
                return ResponseMock(DRIVE_ROOT_WORKING)
            if data[0].get("drivewsid") == "FOLDER::com.apple.CloudDocs::documents":
                return ResponseMock(DRIVE_ROOT_INVALID)
            if (
                data[0].get("drivewsid")
                == "FOLDER::com.apple.CloudDocs::1C7F1760-D940-480F-8C4F-005824A4E05B"
            ):
                return ResponseMock(DRIVE_FOLDER_WORKING)
            if (
                data[0].get("drivewsid")
                == "FOLDER::com.apple.CloudDocs::D5AA0425-E84F-4501-AF5D-60F1D92648CF"
            ):
                return ResponseMock(DRIVE_SUBFOLDER_WORKING)
        # Drive download
        if "com.apple.CloudDocs/download/by_id" in url and method == "GET":
            if params.get("document_id") == "516C896C-6AA5-4A30-B30E-5502C2333DAE":
                return ResponseMock(DRIVE_FILE_DOWNLOAD_WORKING)
        if "icloud-content.com" in url and method == "GET":
            if "Scanned+document+1.pdf" in url:
                return ResponseMock({}, raw=open(".gitignore", "rb"))

        # Find My iPhone
        if "fmi" in url and method == "POST":
            return ResponseMock(FMI_FAMILY_WORKING)

        return None


class PyiCloudServiceMock(base.PyiCloudService):
    """Mocked PyiCloudService."""

    def __init__(
        self,
        apple_id,
        password=None,
        cookie_directory=None,
        verify=True,
        client_id=None,
        with_family=True,
    ):
        """Set up pyicloud service mock."""
        base.PyiCloudSession = PyiCloudSessionMock
        base.PyiCloudService.__init__(
            self, apple_id, password, cookie_directory, verify, client_id, with_family
        )