import json import time from random import randint from typing import Any, Dict, List, Optional, Tuple, Union from urllib.parse import parse_qs from flask import Flask from flask.testing import FlaskClient from urllib3.util import parse_url from werkzeug.test import TestResponse from fittrackee import db from fittrackee.oauth2.client import create_oauth_client from fittrackee.oauth2.models import OAuth2Client, OAuth2Token from fittrackee.users.models import User from .custom_asserts import ( assert_errored_response, assert_oauth_errored_response, ) from .utils import TEST_OAUTH_CLIENT_METADATA, random_email, random_string class RandomMixin: @staticmethod def random_string( length: Optional[int] = None, prefix: Optional[str] = None, suffix: Optional[str] = None, ) -> str: return random_string(length, prefix, suffix) @staticmethod def random_domain() -> str: return random_string(prefix='https://', suffix='com') @staticmethod def random_email() -> str: return random_email() @staticmethod def random_int(min_val: int = 0, max_val: int = 999999) -> int: return randint(min_val, max_val) class OAuth2Mixin(RandomMixin): @staticmethod def create_oauth_client( user: User, metadata: Optional[Dict] = None, scope: Optional[str] = None, ) -> OAuth2Client: client_metadata = ( TEST_OAUTH_CLIENT_METADATA if metadata is None else metadata ) if scope is not None: client_metadata['scope'] = scope oauth_client = create_oauth_client(client_metadata, user) db.session.add(oauth_client) db.session.commit() return oauth_client def create_oauth2_token( self, oauth_client: OAuth2Client, issued_at: Optional[int] = None, access_token_revoked_at: Optional[int] = 0, expires_in: Optional[int] = 1000, ) -> OAuth2Token: issued_at = issued_at if issued_at else int(time.time()) token = OAuth2Token( client_id=oauth_client.client_id, access_token=self.random_string(), refresh_token=self.random_string(), issued_at=issued_at, access_token_revoked_at=access_token_revoked_at, expires_in=expires_in, ) db.session.add(token) db.session.commit() return token class ApiTestCaseMixin(OAuth2Mixin, RandomMixin): @staticmethod def get_test_client_and_auth_token( app: Flask, user_email: str ) -> Tuple[FlaskClient, str]: client = app.test_client() resp_login = client.post( '/api/auth/login', data=json.dumps( dict( email=user_email, password='12345678', ) ), content_type='application/json', ) auth_token = json.loads(resp_login.data.decode())['auth_token'] return client, auth_token @staticmethod def authorize_client( client: FlaskClient, oauth_client: OAuth2Client, auth_token: str, scope: Optional[str] = None, ) -> Union[List[str], str]: response = client.post( '/api/oauth/authorize', data={ 'client_id': oauth_client.client_id, 'confirm': True, 'response_type': 'code', 'scope': 'read' if not scope else scope, }, headers=dict( Authorization=f'Bearer {auth_token}', content_type='multipart/form-data', ), ) data = json.loads(response.data.decode()) parsed_url = parse_url(data['redirect_url']) code = parse_qs(parsed_url.query).get('code', '') return code def create_oauth_client_and_issue_token( self, app: Flask, user: User, scope: Optional[str] = None ) -> Tuple[FlaskClient, OAuth2Client, str, str]: client, auth_token = self.get_test_client_and_auth_token( app, user.email ) oauth_client = self.create_oauth_client(user, scope=scope) code = self.authorize_client( client, oauth_client, auth_token, scope=scope ) response = client.post( '/api/oauth/token', data={ 'client_id': oauth_client.client_id, 'client_secret': oauth_client.client_secret, 'grant_type': 'authorization_code', 'code': code, }, headers=dict(content_type='multipart/form-data'), ) data = json.loads(response.data.decode()) return client, oauth_client, data.get('access_token'), auth_token @staticmethod def assert_400( response: TestResponse, error_message: Optional[str] = 'invalid payload', status: Optional[str] = 'error', ) -> Dict: return assert_errored_response( response, 400, error_message=error_message, status=status ) @staticmethod def assert_401( response: TestResponse, error_message: Optional[str] = 'provide a valid auth token', ) -> Dict: return assert_errored_response( response, 401, error_message=error_message ) @staticmethod def assert_403( response: TestResponse, error_message: Optional[str] = 'you do not have permissions', ) -> Dict: return assert_errored_response(response, 403, error_message) @staticmethod def assert_404(response: TestResponse) -> Dict: return assert_errored_response(response, 404, status='not found') @staticmethod def assert_404_with_entity(response: TestResponse, entity: str) -> Dict: error_message = f'{entity} does not exist' return assert_errored_response( response, 404, error_message=error_message, status='not found' ) @staticmethod def assert_404_with_message( response: TestResponse, error_message: str ) -> Dict: return assert_errored_response( response, 404, error_message=error_message, status='not found' ) @staticmethod def assert_413( response: TestResponse, error_message: Optional[str] = None, match: Optional[str] = None, ) -> Dict: return assert_errored_response( response, 413, error_message=error_message, status='fail', match=match, ) @staticmethod def assert_500( response: TestResponse, error_message: Optional[str] = ( 'error, please try again or contact the administrator' ), status: Optional[str] = 'error', ) -> Dict: return assert_errored_response( response, 500, error_message=error_message, status=status ) @staticmethod def assert_unsupported_grant_type(response: TestResponse) -> Dict: return assert_oauth_errored_response( response, 400, error='unsupported_grant_type' ) @staticmethod def assert_invalid_client(response: TestResponse) -> Dict: return assert_oauth_errored_response( response, 400, error='invalid_client', ) @staticmethod def assert_invalid_request(response: TestResponse) -> Dict: return assert_oauth_errored_response( response, 400, error='invalid_request', ) @staticmethod def assert_invalid_token(response: TestResponse) -> Dict: return assert_oauth_errored_response( response, 401, error='invalid_token', error_description=( 'The access token provided is expired, revoked, malformed, ' 'or invalid for other reasons.' ), ) @staticmethod def assert_insufficient_scope(response: TestResponse) -> Dict: return assert_oauth_errored_response( response, 403, error='insufficient_scope', error_description=( 'The request requires higher privileges than provided by ' 'the access token.' ), ) @staticmethod def assert_not_insufficient_scope_error(response: TestResponse) -> None: assert response.status_code != 403 def assert_response_scope( self, response: TestResponse, can_access: bool ) -> None: if can_access: self.assert_not_insufficient_scope_error(response) else: self.assert_insufficient_scope(response) class CallArgsMixin: @staticmethod def get_args(call_args: Any) -> Any: if len(call_args) == 2: args, _ = call_args else: _, args, _ = call_args return args