<commit_before>from base64 import b64encode
try:
    from urllib.parse import urlencode
except ImportError:
    from urllib import urlencode

from django.test import RequestFactory, override_settings
from django.test import TestCase
from jwkest.jwk import KEYS
from jwkest.jwt import JWT
from mock import patch

from oidc_provider.lib.utils.token import *
from oidc_provider.tests.app.utils import *
from oidc_provider.views import *

class TokenTestCase(TestCase):
    """
    To obtain an Access Token and an ID Token, the RP Client sends a
    Token Request to the Token Endpoint to obtain a Token Response
    when using the Authorization Code Flow.
    """

    def setUp(self):
        self.factory = RequestFactory()
        self.user = create_fake_user()
        self.client = create_fake_client(response_type='code')
        create_rsakey()

    def _auth_code_post_data(self, code):
        """
        All the data that will be POSTed to the Token Endpoint.
        """
        post_data = {
            'client_id': self.client.client_id,
            'client_secret': self.client.client_secret,
            'redirect_uri': self.client.default_redirect_uri,
            'grant_type': 'authorization_code',
            'code': code,
            'state': uuid.uuid4().hex,
        }

        return post_data

    def _refresh_token_post_data(self, refresh_token):
        """
        All the data that will be POSTed to the Token Endpoint.
        """
        post_data = {
            'client_id': self.client.client_id,
            'client_secret': self.client.client_secret,
            'grant_type': 'refresh_token',
            'refresh_token': refresh_token,
        }

        return post_data

    def _post_request(self, post_data, extras={}):
        """
        Makes a request to the token endpoint by sending the
        `post_data` parameters using the 'application/x-www-form-urlencoded'
        format.
        """
        url = reverse('oidc_provider:token')

        request = self.factory.post(url,
            data=urlencode(post_data),
            content_type='application/x-www-form-urlencoded',
            **extras)

        response = TokenView.as_view()(request)

        return response

    def _create_code(self):
        """
        Generate a valid grant code.
        """
        code = create_code(
            user=self.user,
            client=self.client,
            scope=['openid', 'email'],
            nonce=FAKE_NONCE)
        code.save()

        return code

    def _get_keys(self):
        """
        Get public key from discovery.
        """
        request = self.factory.get(reverse('oidc_provider:jwks'))
        response = JwksView.as_view()(request)
        jwks_dic = json.loads(response.content.decode('utf-8'))
        SIGKEYS = KEYS()
        SIGKEYS.load_dict(jwks_dic)
        return SIGKEYS

    def _get_userinfo(self, access_token):
        url = reverse('oidc_provider:userinfo')
        request = self.factory.get(url)
        request.META['HTTP_AUTHORIZATION'] = 'Bearer ' + access_token

        return userinfo(request)

    @override_settings(OIDC_TOKEN_EXPIRE=720)
    def test_authorization_code(self):
        """
        We MUST validate the signature of the ID Token according to JWS
        using the algorithm specified in the alg Header Parameter of
        the JOSE Header.
        """
        SIGKEYS = self._get_keys()
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWS().verify_compact(response_dic['id_token'].encode('utf-8'), SIGKEYS)

        token = Token.objects.get(user=self.user)
        self.assertEqual(response_dic['access_token'], token.access_token)
        self.assertEqual(response_dic['refresh_token'], token.refresh_token)
        self.assertEqual(response_dic['token_type'], 'bearer')
        self.assertEqual(response_dic['expires_in'], 720)
        self.assertEqual(id_token['sub'], str(self.user.id))
        self.assertEqual(id_token['aud'], self.client.client_id)

    def test_refresh_token(self):
        """
        A request to the Token Endpoint can also use a Refresh Token
        by using the grant_type value refresh_token, as described in
        Section 6 of OAuth 2.0 [RFC6749].
        """
        SIGKEYS = self._get_keys()

        # Retrieve refresh token
        code = self._create_code()
        post_data = self._auth_code_post_data(code=code.code)
        real_now = timezone.now
        with patch('oidc_provider.lib.utils.token.timezone.now') as now:
            now.return_value = real_now()
            response = self._post_request(post_data)

        response_dic1 = json.loads(response.content.decode('utf-8'))
        id_token1 = JWS().verify_compact(response_dic1['id_token'].encode('utf-8'), SIGKEYS)

        # Use refresh token to obtain new token
        post_data = self._refresh_token_post_data(response_dic1['refresh_token'])
        with patch('oidc_provider.lib.utils.token.timezone.now') as now:
            now.return_value = real_now() + timedelta(minutes=10)
            response = self._post_request(post_data)

        response_dic2 = json.loads(response.content.decode('utf-8'))
        id_token2 = JWS().verify_compact(response_dic2['id_token'].encode('utf-8'), SIGKEYS)

        self.assertNotEqual(response_dic1['id_token'], response_dic2['id_token'])
        self.assertNotEqual(response_dic1['access_token'], response_dic2['access_token'])
        self.assertNotEqual(response_dic1['refresh_token'], response_dic2['refresh_token'])

        # http://openid.net/specs/openid-connect-core-1_0.html#rfc.section.12.2
        self.assertEqual(id_token1['iss'], id_token2['iss'])
        self.assertEqual(id_token1['sub'], id_token2['sub'])
        self.assertNotEqual(id_token1['iat'], id_token2['iat'])
        self.assertEqual(id_token1['aud'], id_token2['aud'])
        self.assertEqual(id_token1['auth_time'], id_token2['auth_time'])
        self.assertEqual(id_token1.get('azp'), id_token2.get('azp'))

        # Refresh token can't be reused
        post_data = self._refresh_token_post_data(response_dic1['refresh_token'])
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

        # Old access token is invalidated
        self.assertEqual(self._get_userinfo(response_dic1['access_token']).status_code, 401)
        self.assertEqual(self._get_userinfo(response_dic2['access_token']).status_code, 200)

        # Empty refresh token is invalid
        post_data = self._refresh_token_post_data('')
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

        # No refresh token is invalid
        post_data = self._refresh_token_post_data('')
        del post_data['refresh_token']
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

    def test_request_methods(self):
        """
        Client sends an HTTP POST request to the Token Endpoint. Other request
        methods MUST NOT be allowed.
        """
        url = reverse('oidc_provider:token')

        requests = [
            self.factory.get(url),
            self.factory.put(url),
            self.factory.delete(url),
        ]

        for request in requests:
            response = TokenView.as_view()(request)

            self.assertEqual(response.status_code == 405, True,
                msg=request.method+' request does not return a 405 status.')

        request = self.factory.post(url)

        response = TokenView.as_view()(request)

        self.assertEqual(response.status_code == 400, True,
                msg=request.method+' request does not return a 400 status.')

    def test_client_authentication(self):
        """
        The authorization server support including the
        client credentials in the request-body using the `client_id` and
        `client_secret`parameters.

        See: http://tools.ietf.org/html/rfc6749#section-2.3.1
        """
        code = self._create_code()

        # Test a valid request to the token endpoint.
        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                False,
                msg='Client authentication fails using request-body credentials.')

        # Now, test with an invalid client_id.
        invalid_data = post_data.copy()
        invalid_data['client_id'] = self.client.client_id * 2 # Fake id.

        # Create another grant code.
        code = self._create_code()
        invalid_data['code'] = code.code

        response = self._post_request(invalid_data)

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                True,
                msg='Client authentication success with an invalid "client_id".')

        # Now, test using HTTP Basic Authentication method.
        basicauth_data = post_data.copy()

        # Create another grant code.
        code = self._create_code()
        basicauth_data['code'] = code.code

        del basicauth_data['client_id']
        del basicauth_data['client_secret']

        # Generate HTTP Basic Auth header with id and secret.
        user_pass = self.client.client_id + ':' + self.client.client_secret
        auth_header = b'Basic ' + b64encode(user_pass.encode('utf-8'))
        response = self._post_request(basicauth_data, {
                       'HTTP_AUTHORIZATION': auth_header.decode('utf-8'),
                   })
        response.content.decode('utf-8')

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                False,
                msg='Client authentication fails using HTTP Basic Auth.')

    def test_access_token_contains_nonce(self):
        """
        If present in the Authentication Request, Authorization Servers MUST
        include a nonce Claim in the ID Token with the Claim Value being the
        nonce value sent in the Authentication Request.
        If the client does not supply a nonce parameter, it SHOULD not be
        included in the `id_token`.

        See http://openid.net/specs/openid-connect-core-1_0.html#IDToken
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('nonce'), FAKE_NONCE)

        # Client does not supply a nonce parameter.
        code.nonce = ''
        code.save()

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('nonce'), None)

    def test_idtoken_sign_validation(self):
        """
        We MUST validate the signature of the ID Token according to JWS
        using the algorithm specified in the alg Header Parameter of
        the JOSE Header.
        """
        SIGKEYS = self._get_keys()
        RSAKEYS = [ k for k in SIGKEYS if k.kty == 'RSA' ]

        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWS().verify_compact(response_dic['id_token'].encode('utf-8'), RSAKEYS)

    @override_settings(OIDC_IDTOKEN_SUB_GENERATOR='oidc_provider.tests.app.utils.fake_sub_generator')
    def test_custom_sub_generator(self):
        """
        Test custom function for setting OIDC_IDTOKEN_SUB_GENERATOR.
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('sub'), self.user.email)

    @override_settings(OIDC_IDTOKEN_PROCESSING_HOOK='oidc_provider.tests.app.utils.fake_idtoken_processing_hook')
    def test_additional_idtoken_processing_hook(self):
        """
        Test custom function for setting OIDC_IDTOKEN_PROCESSING_HOOK.
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('test_idtoken_processing_hook'), FAKE_RANDOM_STRING)
<commit_msg>added 'user' into default_idtoken_processing_hook<commit_after>from base64 import b64encode
try:
    from urllib.parse import urlencode
except ImportError:
    from urllib import urlencode

from django.test import RequestFactory, override_settings
from django.test import TestCase
from jwkest.jwk import KEYS
from jwkest.jwt import JWT
from mock import patch

from oidc_provider.lib.utils.token import *
from oidc_provider.tests.app.utils import *
from oidc_provider.views import *

class TokenTestCase(TestCase):
    """
    To obtain an Access Token and an ID Token, the RP Client sends a
    Token Request to the Token Endpoint to obtain a Token Response
    when using the Authorization Code Flow.
    """

    def setUp(self):
        self.factory = RequestFactory()
        self.user = create_fake_user()
        self.client = create_fake_client(response_type='code')
        create_rsakey()

    def _auth_code_post_data(self, code):
        """
        All the data that will be POSTed to the Token Endpoint.
        """
        post_data = {
            'client_id': self.client.client_id,
            'client_secret': self.client.client_secret,
            'redirect_uri': self.client.default_redirect_uri,
            'grant_type': 'authorization_code',
            'code': code,
            'state': uuid.uuid4().hex,
        }

        return post_data

    def _refresh_token_post_data(self, refresh_token):
        """
        All the data that will be POSTed to the Token Endpoint.
        """
        post_data = {
            'client_id': self.client.client_id,
            'client_secret': self.client.client_secret,
            'grant_type': 'refresh_token',
            'refresh_token': refresh_token,
        }

        return post_data

    def _post_request(self, post_data, extras={}):
        """
        Makes a request to the token endpoint by sending the
        `post_data` parameters using the 'application/x-www-form-urlencoded'
        format.
        """
        url = reverse('oidc_provider:token')

        request = self.factory.post(url,
            data=urlencode(post_data),
            content_type='application/x-www-form-urlencoded',
            **extras)

        response = TokenView.as_view()(request)

        return response

    def _create_code(self):
        """
        Generate a valid grant code.
        """
        code = create_code(
            user=self.user,
            client=self.client,
            scope=['openid', 'email'],
            nonce=FAKE_NONCE)
        code.save()

        return code

    def _get_keys(self):
        """
        Get public key from discovery.
        """
        request = self.factory.get(reverse('oidc_provider:jwks'))
        response = JwksView.as_view()(request)
        jwks_dic = json.loads(response.content.decode('utf-8'))
        SIGKEYS = KEYS()
        SIGKEYS.load_dict(jwks_dic)
        return SIGKEYS

    def _get_userinfo(self, access_token):
        url = reverse('oidc_provider:userinfo')
        request = self.factory.get(url)
        request.META['HTTP_AUTHORIZATION'] = 'Bearer ' + access_token

        return userinfo(request)

    @override_settings(OIDC_TOKEN_EXPIRE=720)
    def test_authorization_code(self):
        """
        We MUST validate the signature of the ID Token according to JWS
        using the algorithm specified in the alg Header Parameter of
        the JOSE Header.
        """
        SIGKEYS = self._get_keys()
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWS().verify_compact(response_dic['id_token'].encode('utf-8'), SIGKEYS)

        token = Token.objects.get(user=self.user)
        self.assertEqual(response_dic['access_token'], token.access_token)
        self.assertEqual(response_dic['refresh_token'], token.refresh_token)
        self.assertEqual(response_dic['token_type'], 'bearer')
        self.assertEqual(response_dic['expires_in'], 720)
        self.assertEqual(id_token['sub'], str(self.user.id))
        self.assertEqual(id_token['aud'], self.client.client_id)

    def test_refresh_token(self):
        """
        A request to the Token Endpoint can also use a Refresh Token
        by using the grant_type value refresh_token, as described in
        Section 6 of OAuth 2.0 [RFC6749].
        """
        SIGKEYS = self._get_keys()

        # Retrieve refresh token
        code = self._create_code()
        post_data = self._auth_code_post_data(code=code.code)
        real_now = timezone.now
        with patch('oidc_provider.lib.utils.token.timezone.now') as now:
            now.return_value = real_now()
            response = self._post_request(post_data)

        response_dic1 = json.loads(response.content.decode('utf-8'))
        id_token1 = JWS().verify_compact(response_dic1['id_token'].encode('utf-8'), SIGKEYS)

        # Use refresh token to obtain new token
        post_data = self._refresh_token_post_data(response_dic1['refresh_token'])
        with patch('oidc_provider.lib.utils.token.timezone.now') as now:
            now.return_value = real_now() + timedelta(minutes=10)
            response = self._post_request(post_data)

        response_dic2 = json.loads(response.content.decode('utf-8'))
        id_token2 = JWS().verify_compact(response_dic2['id_token'].encode('utf-8'), SIGKEYS)

        self.assertNotEqual(response_dic1['id_token'], response_dic2['id_token'])
        self.assertNotEqual(response_dic1['access_token'], response_dic2['access_token'])
        self.assertNotEqual(response_dic1['refresh_token'], response_dic2['refresh_token'])

        # http://openid.net/specs/openid-connect-core-1_0.html#rfc.section.12.2
        self.assertEqual(id_token1['iss'], id_token2['iss'])
        self.assertEqual(id_token1['sub'], id_token2['sub'])
        self.assertNotEqual(id_token1['iat'], id_token2['iat'])
        self.assertEqual(id_token1['aud'], id_token2['aud'])
        self.assertEqual(id_token1['auth_time'], id_token2['auth_time'])
        self.assertEqual(id_token1.get('azp'), id_token2.get('azp'))

        # Refresh token can't be reused
        post_data = self._refresh_token_post_data(response_dic1['refresh_token'])
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

        # Old access token is invalidated
        self.assertEqual(self._get_userinfo(response_dic1['access_token']).status_code, 401)
        self.assertEqual(self._get_userinfo(response_dic2['access_token']).status_code, 200)

        # Empty refresh token is invalid
        post_data = self._refresh_token_post_data('')
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

        # No refresh token is invalid
        post_data = self._refresh_token_post_data('')
        del post_data['refresh_token']
        response = self._post_request(post_data)
        self.assertIn('invalid_grant', response.content.decode('utf-8'))

    def test_request_methods(self):
        """
        Client sends an HTTP POST request to the Token Endpoint. Other request
        methods MUST NOT be allowed.
        """
        url = reverse('oidc_provider:token')

        requests = [
            self.factory.get(url),
            self.factory.put(url),
            self.factory.delete(url),
        ]

        for request in requests:
            response = TokenView.as_view()(request)

            self.assertEqual(response.status_code == 405, True,
                msg=request.method+' request does not return a 405 status.')

        request = self.factory.post(url)

        response = TokenView.as_view()(request)

        self.assertEqual(response.status_code == 400, True,
                msg=request.method+' request does not return a 400 status.')

    def test_client_authentication(self):
        """
        The authorization server support including the
        client credentials in the request-body using the `client_id` and
        `client_secret`parameters.

        See: http://tools.ietf.org/html/rfc6749#section-2.3.1
        """
        code = self._create_code()

        # Test a valid request to the token endpoint.
        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                False,
                msg='Client authentication fails using request-body credentials.')

        # Now, test with an invalid client_id.
        invalid_data = post_data.copy()
        invalid_data['client_id'] = self.client.client_id * 2 # Fake id.

        # Create another grant code.
        code = self._create_code()
        invalid_data['code'] = code.code

        response = self._post_request(invalid_data)

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                True,
                msg='Client authentication success with an invalid "client_id".')

        # Now, test using HTTP Basic Authentication method.
        basicauth_data = post_data.copy()

        # Create another grant code.
        code = self._create_code()
        basicauth_data['code'] = code.code

        del basicauth_data['client_id']
        del basicauth_data['client_secret']

        # Generate HTTP Basic Auth header with id and secret.
        user_pass = self.client.client_id + ':' + self.client.client_secret
        auth_header = b'Basic ' + b64encode(user_pass.encode('utf-8'))
        response = self._post_request(basicauth_data, {
                       'HTTP_AUTHORIZATION': auth_header.decode('utf-8'),
                   })
        response.content.decode('utf-8')

        self.assertEqual('invalid_client' in response.content.decode('utf-8'),
                False,
                msg='Client authentication fails using HTTP Basic Auth.')

    def test_access_token_contains_nonce(self):
        """
        If present in the Authentication Request, Authorization Servers MUST
        include a nonce Claim in the ID Token with the Claim Value being the
        nonce value sent in the Authentication Request.
        If the client does not supply a nonce parameter, it SHOULD not be
        included in the `id_token`.

        See http://openid.net/specs/openid-connect-core-1_0.html#IDToken
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('nonce'), FAKE_NONCE)

        # Client does not supply a nonce parameter.
        code.nonce = ''
        code.save()

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('nonce'), None)

    def test_idtoken_sign_validation(self):
        """
        We MUST validate the signature of the ID Token according to JWS
        using the algorithm specified in the alg Header Parameter of
        the JOSE Header.
        """
        SIGKEYS = self._get_keys()
        RSAKEYS = [ k for k in SIGKEYS if k.kty == 'RSA' ]

        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)
        response_dic = json.loads(response.content.decode('utf-8'))

        id_token = JWS().verify_compact(response_dic['id_token'].encode('utf-8'), RSAKEYS)

    @override_settings(OIDC_IDTOKEN_SUB_GENERATOR='oidc_provider.tests.app.utils.fake_sub_generator')
    def test_custom_sub_generator(self):
        """
        Test custom function for setting OIDC_IDTOKEN_SUB_GENERATOR.
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('sub'), self.user.email)

    @override_settings(OIDC_IDTOKEN_PROCESSING_HOOK='oidc_provider.tests.app.utils.fake_idtoken_processing_hook')
    def test_additional_idtoken_processing_hook(self):
        """
        Test custom function for setting OIDC_IDTOKEN_PROCESSING_HOOK.
        """
        code = self._create_code()

        post_data = self._auth_code_post_data(code=code.code)

        response = self._post_request(post_data)

        response_dic = json.loads(response.content.decode('utf-8'))
        id_token = JWT().unpack(response_dic['id_token'].encode('utf-8')).payload()

        self.assertEqual(id_token.get('test_idtoken_processing_hook'), FAKE_RANDOM_STRING)
        self.assertEqual(id_token.get('test_idtoken_processing_hook_user_email'), self.user.email)
<|endoftext|>


<commit_before>"""Support for the Italian train system using ViaggiaTreno API."""
from __future__ import annotations

import asyncio
from http import HTTPStatus
import logging
import time

import aiohttp
import async_timeout
import voluptuous as vol

from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
from homeassistant.const import ATTR_ATTRIBUTION, TIME_MINUTES
from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_get_clientsession
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType

_LOGGER = logging.getLogger(__name__)

ATTRIBUTION = "Powered by ViaggiaTreno Data"

VIAGGIATRENO_ENDPOINT = (
	"http://www.viaggiatreno.it/infomobilita/"
    "resteasy/viaggiatreno/andamentoTreno/"
    "{station_id}/{train_id}/{timestamp}"
)

REQUEST_TIMEOUT = 5  # seconds
ICON = "mdi:train"
MONITORED_INFO = [
    "categoria",
    "compOrarioArrivoZeroEffettivo",
    "compOrarioPartenzaZeroEffettivo",
    "destinazione",
    "numeroTreno",
    "orarioArrivo",
    "orarioPartenza",
    "origine",
    "subTitle",
]

DEFAULT_NAME = "Train {}"

CONF_NAME = "train_name"
CONF_STATION_ID = "station_id"
CONF_STATION_NAME = "station_name"
CONF_TRAIN_ID = "train_id"

ARRIVED_STRING = "Arrived"
CANCELLED_STRING = "Cancelled"
NOT_DEPARTED_STRING = "Not departed yet"
NO_INFORMATION_STRING = "No information for this train now"

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_TRAIN_ID): cv.string,
        vol.Required(CONF_STATION_ID): cv.string,
        vol.Optional(CONF_NAME): cv.string,
    }
)

async def async_setup_platform(
    hass: HomeAssistant,
    config: ConfigType,
    async_add_entities: AddEntitiesCallback,
    discovery_info: DiscoveryInfoType | None = None,
) -> None:
    """Set up the ViaggiaTreno platform."""
    train_id = config.get(CONF_TRAIN_ID)
    station_id = config.get(CONF_STATION_ID)
    if not (name := config.get(CONF_NAME)):
        name = DEFAULT_NAME.format(train_id)
    async_add_entities([ViaggiaTrenoSensor(train_id, station_id, name)])

async def async_http_request(hass, uri):
    """Perform actual request."""
    try:
        session = async_get_clientsession(hass)
        async with async_timeout.timeout(REQUEST_TIMEOUT):
            req = await session.get(uri)
        if req.status != HTTPStatus.OK:
            return {"error": req.status}
        json_response = await req.json()
        return json_response
    except (asyncio.TimeoutError, aiohttp.ClientError) as exc:
        _LOGGER.error("Cannot connect to ViaggiaTreno API endpoint: %s", exc)
    except ValueError:
        _LOGGER.error("Received non-JSON data from ViaggiaTreno API endpoint")

class ViaggiaTrenoSensor(SensorEntity):
    """Implementation of a ViaggiaTreno sensor."""

    def __init__(self, train_id, station_id, name):
        """Initialize the sensor."""
        self._state = None
        self._attributes = {}
        self._unit = ""
        self._icon = ICON
        self._station_id = station_id
        self._name = name

        self.uri = VIAGGIATRENO_ENDPOINT.format(
            station_id=station_id, train_id=train_id, timestamp=int(time.time()) * 1000
        )

    @property
    def name(self):
        """Return the name of the sensor."""
        return self._name

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self._state

    @property
    def icon(self):
        """Icon to use in the frontend, if any."""
        return self._icon

    @property
    def native_unit_of_measurement(self):
        """Return the unit of measurement."""
        return self._unit

    @property
    def extra_state_attributes(self):
        """Return extra attributes."""
        self._attributes[ATTR_ATTRIBUTION] = ATTRIBUTION
        return self._attributes

    @staticmethod
    def has_departed(data):
        """Check if the train has actually departed."""
        try:
            first_station = data["fermate"][0]
            if data["oraUltimoRilevamento"] or first_station["effettiva"]:
                return True
        except ValueError:
            _LOGGER.error("Cannot fetch first station: %s", data)
        return False

    @staticmethod
    def has_arrived(data):
        """Check if the train has already arrived."""
        last_station = data["fermate"][-1]
        if not last_station["effettiva"]:
            return False
        return True

    @staticmethod
    def is_cancelled(data):
        """Check if the train is cancelled."""
        if data["tipoTreno"] == "ST" and data["provvedimento"] == 1:
            return True
        return False

    async def async_update(self):
        """Update state."""
        uri = self.uri
        res = await async_http_request(self.hass, uri)
        if res.get("error", ""):
            if res["error"] == 204:
                self._state = NO_INFORMATION_STRING
                self._unit = ""
            else:
                self._state = "Error: {}".format(res["error"])
                self._unit = ""
        else:
            for i in MONITORED_INFO:
                self._attributes[i] = res[i]

            if self.is_cancelled(res):
                self._state = CANCELLED_STRING
                self._icon = "mdi:cancel"
                self._unit = ""
            elif not self.has_departed(res):
                self._state = NOT_DEPARTED_STRING
                self._unit = ""
            elif self.has_arrived(res):
                self._state = ARRIVED_STRING
                self._unit = ""
            else:
                self._state = res.get("ritardo")
                self._unit = TIME_MINUTES
                self._icon = ICON
<commit_msg>Format viaggiatreno/sensor.py (#74161)<commit_after>"""Support for the Italian train system using ViaggiaTreno API."""
from __future__ import annotations

import asyncio
from http import HTTPStatus
import logging
import time

import aiohttp
import async_timeout
import voluptuous as vol

from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
from homeassistant.const import ATTR_ATTRIBUTION, TIME_MINUTES
from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_get_clientsession
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType

_LOGGER = logging.getLogger(__name__)

ATTRIBUTION = "Powered by ViaggiaTreno Data"

VIAGGIATRENO_ENDPOINT = (
    "http://www.viaggiatreno.it/infomobilita/"
    "resteasy/viaggiatreno/andamentoTreno/"
    "{station_id}/{train_id}/{timestamp}"
)

REQUEST_TIMEOUT = 5  # seconds
ICON = "mdi:train"
MONITORED_INFO = [
    "categoria",
    "compOrarioArrivoZeroEffettivo",
    "compOrarioPartenzaZeroEffettivo",
    "destinazione",
    "numeroTreno",
    "orarioArrivo",
    "orarioPartenza",
    "origine",
    "subTitle",
]

DEFAULT_NAME = "Train {}"

CONF_NAME = "train_name"
CONF_STATION_ID = "station_id"
CONF_STATION_NAME = "station_name"
CONF_TRAIN_ID = "train_id"

ARRIVED_STRING = "Arrived"
CANCELLED_STRING = "Cancelled"
NOT_DEPARTED_STRING = "Not departed yet"
NO_INFORMATION_STRING = "No information for this train now"

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_TRAIN_ID): cv.string,
        vol.Required(CONF_STATION_ID): cv.string,
        vol.Optional(CONF_NAME): cv.string,
    }
)

async def async_setup_platform(
    hass: HomeAssistant,
    config: ConfigType,
    async_add_entities: AddEntitiesCallback,
    discovery_info: DiscoveryInfoType | None = None,
) -> None:
    """Set up the ViaggiaTreno platform."""
    train_id = config.get(CONF_TRAIN_ID)
    station_id = config.get(CONF_STATION_ID)
    if not (name := config.get(CONF_NAME)):
        name = DEFAULT_NAME.format(train_id)
    async_add_entities([ViaggiaTrenoSensor(train_id, station_id, name)])

async def async_http_request(hass, uri):
    """Perform actual request."""
    try:
        session = async_get_clientsession(hass)
        async with async_timeout.timeout(REQUEST_TIMEOUT):
            req = await session.get(uri)
        if req.status != HTTPStatus.OK:
            return {"error": req.status}
        json_response = await req.json()
        return json_response
    except (asyncio.TimeoutError, aiohttp.ClientError) as exc:
        _LOGGER.error("Cannot connect to ViaggiaTreno API endpoint: %s", exc)
    except ValueError:
        _LOGGER.error("Received non-JSON data from ViaggiaTreno API endpoint")

class ViaggiaTrenoSensor(SensorEntity):
    """Implementation of a ViaggiaTreno sensor."""

    def __init__(self, train_id, station_id, name):
        """Initialize the sensor."""
        self._state = None
        self._attributes = {}
        self._unit = ""
        self._icon = ICON
        self._station_id = station_id
        self._name = name

        self.uri = VIAGGIATRENO_ENDPOINT.format(
            station_id=station_id, train_id=train_id, timestamp=int(time.time()) * 1000
        )

    @property
    def name(self):
        """Return the name of the sensor."""
        return self._name

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self._state

    @property
    def icon(self):
        """Icon to use in the frontend, if any."""
        return self._icon

    @property
    def native_unit_of_measurement(self):
        """Return the unit of measurement."""
        return self._unit

    @property
    def extra_state_attributes(self):
        """Return extra attributes."""
        self._attributes[ATTR_ATTRIBUTION] = ATTRIBUTION
        return self._attributes

    @staticmethod
    def has_departed(data):
        """Check if the train has actually departed."""
        try:
            first_station = data["fermate"][0]
            if data["oraUltimoRilevamento"] or first_station["effettiva"]:
                return True
        except ValueError:
            _LOGGER.error("Cannot fetch first station: %s", data)
        return False

    @staticmethod
    def has_arrived(data):
        """Check if the train has already arrived."""
        last_station = data["fermate"][-1]
        if not last_station["effettiva"]:
            return False
        return True

    @staticmethod
    def is_cancelled(data):
        """Check if the train is cancelled."""
        if data["tipoTreno"] == "ST" and data["provvedimento"] == 1:
            return True
        return False

    async def async_update(self):
        """Update state."""
        uri = self.uri
        res = await async_http_request(self.hass, uri)
        if res.get("error", ""):
            if res["error"] == 204:
                self._state = NO_INFORMATION_STRING
                self._unit = ""
            else:
                self._state = "Error: {}".format(res["error"])
                self._unit = ""
        else:
            for i in MONITORED_INFO:
                self._attributes[i] = res[i]

            if self.is_cancelled(res):
                self._state = CANCELLED_STRING
                self._icon = "mdi:cancel"
                self._unit = ""
            elif not self.has_departed(res):
                self._state = NOT_DEPARTED_STRING
                self._unit = ""
            elif self.has_arrived(res):
                self._state = ARRIVED_STRING
                self._unit = ""
            else:
                self._state = res.get("ritardo")
                self._unit = TIME_MINUTES
                self._icon = ICON
<|endoftext|>


<commit_before>from selenium import webdriver
from selenium.common.exceptions import WebDriverException, StaleElementReferenceException
from selenium.common.exceptions import NoSuchElementException
from selenium.webdriver.common.by import By
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

from mailpile.httpd import HttpWorker
from tests import MailPileUnittest, get_shared_mailpile

from mailpile.safe_popen import MakePopenUnsafe

MakePopenUnsafe()

class ElementHasClass(object):
    def __init__(self, locator_tuple, class_name):
        self.locator = locator_tuple
        self.class_name = class_name

    def __call__(self, driver):
        try:
            e = driver.find_element(self.locator[0], self.locator[1])
            return self.class_name in e.get_attribute('class')
        except (NoSuchElementException, StaleElementReferenceException):
            return False

class ElementHasNotClass(object):
    def __init__(self, locator_tuple, class_name):
        self.locator = locator_tuple
        self.class_name = class_name

    def __call__(self, driver):
        try:
            e = driver.find_element(self.locator[0], self.locator[1])
            return self.class_name not in e.get_attribute('class')
        except (NoSuchElementException, StaleElementReferenceException):
            return True

class SeleniumScreenshotOnExceptionAspecter(type):
    """Wraps all methods starting with *test* with a screenshot aspect.

      The screenshot file is named *methodname*_screenshot.png.

      Notes:
        This class defines a type that has to be used as a metaclass:

         >>> class Foobar()
         ...    __metaclass__ = SeleniumScreenshotOnExceptionAspecter
         ...
         ...    def take_screenshot(self, filename):
         ...        # take screenshot
         ...        pass

         The class has to provide a take_screenshot(filename) method

      Attributes:
        none
    """

    def __new__(mcs, name, bases, dict):
        for key, value in dict.items():
            if (hasattr(value, "__call__")
                and key != "__metaclass__"
                and key.startswith('test')):
                dict[key] = SeleniumScreenshotOnExceptionAspecter.wrap_method(
                    value)
        return super(SeleniumScreenshotOnExceptionAspecter,
                     mcs).__new__(mcs, name, bases, dict)

    @classmethod
    def wrap_method(mcs, method):
        """Wraps method with a screenshot on exception aspect."""
        # method name has to start with test, otherwise unittest runner
        # won't detect it
        def test_call_wrapper_method(*args, **kw):
            """The wrapper method

              Notes:
                The method name has to start with *test*, otherwise the
                unittest runner won't detect is as a test method

              Args:
                *args: Variable argument list of original method
                **kw: Arbitrary keyword arguments of the original method

              Returns:
                The result of the original method call
            """
            try:
                results = method(*args, **kw)
            except:
                test_self = args[0]
                filename = '%s_screenshot.png' % method.__name__
                test_self.take_screenshot(filename)
                raise

            return results

        return test_call_wrapper_method

class MailpileSeleniumTest(MailPileUnittest):
    """Base class for all selenium GUI tests

        Attributes:
            DRIVER (WebDriver): The webdriver instance

        Examples:

        >>> class Sometest(MailpileSeleniumTest):
        ...
        ...     def test_something(self):
        ...         self.go_to_mailpile_home()
        ...         self.take_screenshot('screen.png')
        ...         self.dump_source_to('source.html')
        ...
        ...         self.navigate_to('Contacts')
        ...
        ...         self.driver.save_screenshot('screen2.png')
        ...         self.assertIn('Contacts', self.driver.title)
    """
    __metaclass__ = SeleniumScreenshotOnExceptionAspecter

    DRIVER = None
    http_worker = None

    def __init__(self, *args, **kwargs):
        MailPileUnittest.__init__(self, *args, **kwargs)

    def setUp(self):
        self.driver = MailpileSeleniumTest.DRIVER

    def tearDown(self):
        #        try:
        #            self.driver.close()
        #        except WebDriverException:
        #            pass
        pass

    @classmethod
    def _get_mailpile_sspec(cls):
        (_, _, config, _) = get_shared_mailpile()
        return (config.sys.http_host, config.sys.http_port)

    @classmethod
    def _get_mailpile_url(cls):
        return 'http://%s:%s' % cls._get_mailpile_sspec()

    @classmethod
    def _start_web_server(cls):
        if not MailpileSeleniumTest.http_worker:
            (mp, session, config, _) = get_shared_mailpile()
            sspec = MailpileSeleniumTest._get_mailpile_sspec()
            MailpileSeleniumTest.http_worker = config.http_worker = HttpWorker(session, sspec)
            config.http_worker.start()

    @classmethod
    def _start_selenium_driver(cls):
        if not MailpileSeleniumTest.DRIVER:
            driver = webdriver.PhantomJS()  # or add to your PATH
            driver.set_window_size(1280, 1024)  # optional
            driver.implicitly_wait(5)
            driver.set_page_load_timeout(5)
            MailpileSeleniumTest.DRIVER = driver

    @classmethod
    def _stop_selenium_driver(cls):
        if MailpileSeleniumTest.DRIVER:
            try:
                MailpileSeleniumTest.DRIVER.quit()
                MailpileSeleniumTest.DRIVER = None
            except WebDriverException:
                pass

    @classmethod
    def setUpClass(cls):
        return  # FIXME: Test disabled

        MailpileSeleniumTest._start_selenium_driver()
        MailpileSeleniumTest._start_web_server()

    @classmethod
    def _stop_web_server(cls):
        if MailpileSeleniumTest.http_worker:
            (mp, _, config, _) = get_shared_mailpile()
            mp._config.http_worker = None
            MailpileSeleniumTest.http_worker.quit()
            MailpileSeleniumTest.http_worker = MailpileSeleniumTest.http_worker = None

    @classmethod
    def tearDownClass(cls):
        return  # FIXME: Test disabled

        MailpileSeleniumTest._stop_web_server()
        MailpileSeleniumTest._stop_selenium_driver()

    def go_to_mailpile_home(self):
        self.driver.get('%s/in/inbox' % self._get_mailpile_url())

    def take_screenshot(self, filename):
        try:
            self.driver.save_screenshot(filename)  # save a screenshot to disk
        except WebDriverException:
            pass

    def dump_source_to(self, filename):
        with open(filename, 'w') as out:
            out.write(self.driver.page_source.encode('utf8'))

    def navigate_to(self, name):
        contacts = self.find_element_by_xpath(
            '//a[@alt="%s"]/span' % name)
        self.assertTrue(contacts.is_displayed())
        contacts.click()

    def submit_form(self, form_id):
        form = self.driver.find_element_by_id(form_id)
        form.submit()

    def fill_form_field(self, field, text):
        input_field = self.driver.find_element_by_name(field)
        input_field.send_keys(text)

    def assert_link_with_text(self, text):
        try:
            self.driver.find_element_by_link_text(text)
        except NoSuchElementException:
            raise AssertionError

    def click_element_with_link_text(self, text):
        try:
            self.driver.find_element_by_link_text(text).click()
        except NoSuchElementException:
            raise AssertionError

    def click_element_with_id(self, element_id):
        self.driver.find_element_by_id(element_id).click()

    def click_element_with_class(self, class_name):
        self.driver.find_element_by_class_name(class_name).click()

    def page_title(self):
        return self.driver.title

    def find_element_by_id(self, id):
        return self.driver.find_element_by_id(id)

    def find_element_containing_text(self, text):
        return self.driver.find_element_by_xpath("//*[contains(.,'%s')]" % text)

    def find_element_by_xpath(self, xpath):
        return self.driver.find_element_by_xpath(xpath)

    def find_element_by_class_name(self, class_name):
        return self.driver.find_element_by_class_name(class_name)

    def assert_text(self, text):
        self.find_element_containing_text(text)

    def wait_until_element_is_visible(self, element_id):
        self.wait_until_element_is_visible_by_locator((By.ID, element_id))

    def wait_until_element_is_visible_by_locator(self, locator_tuple):
        wait = WebDriverWait(self.driver, 10)
        wait.until(EC.visibility_of_element_located(locator_tuple))

    def wait_until_element_is_invisible_by_locator(self, locator_tuple):
        wait = WebDriverWait(self.driver, 10)
        wait.until(EC.invisibility_of_element_located(locator_tuple))

    def wait_until_element_has_class(self, locator_tuple, class_name):
        self.wait_for_element_condition(ElementHasClass(locator_tuple, class_name))

    def wait_until_element_has_not_class(self, locator_tuple, class_name):
        self.wait_for_element_condition(ElementHasNotClass(locator_tuple, class_name))

    def wait_for_element_condition(self, expected_conditions):
        wait = WebDriverWait(self.driver, 10)
        wait.until(expected_conditions)
<commit_msg>More selenium disabling<commit_after>try:
    from selenium import webdriver
    from selenium.common.exceptions import WebDriverException, StaleElementReferenceException
    from selenium.common.exceptions import NoSuchElementException
    from selenium.webdriver.common.by import By
    from selenium.webdriver.support.wait import WebDriverWait
    from selenium.webdriver.support import expected_conditions as EC
except ImportError:
    pass

from mailpile.httpd import HttpWorker
from tests import MailPileUnittest, get_shared_mailpile

from mailpile.safe_popen import MakePopenUnsafe

MakePopenUnsafe()

class ElementHasClass(object):
    def __init__(self, locator_tuple, class_name):
        self.locator = locator_tuple
        self.class_name = class_name

    def __call__(self, driver):
        try:
            e = driver.find_element(self.locator[0], self.locator[1])
            return self.class_name in e.get_attribute('class')
        except (NoSuchElementException, StaleElementReferenceException):
            return False

class ElementHasNotClass(object):
    def __init__(self, locator_tuple, class_name):
        self.locator = locator_tuple
        self.class_name = class_name

    def __call__(self, driver):
        try:
            e = driver.find_element(self.locator[0], self.locator[1])
            return self.class_name not in e.get_attribute('class')
        except (NoSuchElementException, StaleElementReferenceException):
            return True

class SeleniumScreenshotOnExceptionAspecter(type):
    """Wraps all methods starting with *test* with a screenshot aspect.

      The screenshot file is named *methodname*_screenshot.png.

      Notes:
        This class defines a type that has to be used as a metaclass:

         >>> class Foobar()
         ...    __metaclass__ = SeleniumScreenshotOnExceptionAspecter
         ...
         ...    def take_screenshot(self, filename):
         ...        # take screenshot
         ...        pass

         The class has to provide a take_screenshot(filename) method

      Attributes:
        none
    """

    def __new__(mcs, name, bases, dict):
        for key, value in dict.items():
            if (hasattr(value, "__call__")
                and key != "__metaclass__"
                and key.startswith('test')):
                dict[key] = SeleniumScreenshotOnExceptionAspecter.wrap_method(
                    value)
        return super(SeleniumScreenshotOnExceptionAspecter,
                     mcs).__new__(mcs, name, bases, dict)

    @classmethod
    def wrap_method(mcs, method):
        """Wraps method with a screenshot on exception aspect."""
        # method name has to start with test, otherwise unittest runner
        # won't detect it
        def test_call_wrapper_method(*args, **kw):
            """The wrapper method

              Notes:
                The method name has to start with *test*, otherwise the
                unittest runner won't detect is as a test method

              Args:
                *args: Variable argument list of original method
                **kw: Arbitrary keyword arguments of the original method

              Returns:
                The result of the original method call
            """
            try:
                results = method(*args, **kw)
            except:
                test_self = args[0]
                filename = '%s_screenshot.png' % method.__name__
                test_self.take_screenshot(filename)
                raise

            return results

        return test_call_wrapper_method

class MailpileSeleniumTest(MailPileUnittest):
    """Base class for all selenium GUI tests

        Attributes:
            DRIVER (WebDriver): The webdriver instance

        Examples:

        >>> class Sometest(MailpileSeleniumTest):
        ...
        ...     def test_something(self):
        ...         self.go_to_mailpile_home()
        ...         self.take_screenshot('screen.png')
        ...         self.dump_source_to('source.html')
        ...
        ...         self.navigate_to('Contacts')
        ...
        ...         self.driver.save_screenshot('screen2.png')
        ...         self.assertIn('Contacts', self.driver.title)
    """
    __metaclass__ = SeleniumScreenshotOnExceptionAspecter

    DRIVER = None
    http_worker = None

    def __init__(self, *args, **kwargs):
        MailPileUnittest.__init__(self, *args, **kwargs)

    def setUp(self):
        self.driver = MailpileSeleniumTest.DRIVER

    def tearDown(self):
        #        try:
        #            self.driver.close()
        #        except WebDriverException:
        #            pass
        pass

    @classmethod
    def _get_mailpile_sspec(cls):
        (_, _, config, _) = get_shared_mailpile()
        return (config.sys.http_host, config.sys.http_port)

    @classmethod
    def _get_mailpile_url(cls):
        return 'http://%s:%s' % cls._get_mailpile_sspec()

    @classmethod
    def _start_web_server(cls):
        if not MailpileSeleniumTest.http_worker:
            (mp, session, config, _) = get_shared_mailpile()
            sspec = MailpileSeleniumTest._get_mailpile_sspec()
            MailpileSeleniumTest.http_worker = config.http_worker = HttpWorker(session, sspec)
            config.http_worker.start()

    @classmethod
    def _start_selenium_driver(cls):
        if not MailpileSeleniumTest.DRIVER:
            driver = webdriver.PhantomJS()  # or add to your PATH
            driver.set_window_size(1280, 1024)  # optional
            driver.implicitly_wait(5)
            driver.set_page_load_timeout(5)
            MailpileSeleniumTest.DRIVER = driver

    @classmethod
    def _stop_selenium_driver(cls):
        if MailpileSeleniumTest.DRIVER:
            try:
                MailpileSeleniumTest.DRIVER.quit()
                MailpileSeleniumTest.DRIVER = None
            except WebDriverException:
                pass

    @classmethod
    def setUpClass(cls):
        return  # FIXME: Test disabled

        MailpileSeleniumTest._start_selenium_driver()
        MailpileSeleniumTest._start_web_server()

    @classmethod
    def _stop_web_server(cls):
        if MailpileSeleniumTest.http_worker:
            (mp, _, config, _) = get_shared_mailpile()
            mp._config.http_worker = None
            MailpileSeleniumTest.http_worker.quit()
            MailpileSeleniumTest.http_worker = MailpileSeleniumTest.http_worker = None

    @classmethod
    def tearDownClass(cls):
        return  # FIXME: Test disabled

        MailpileSeleniumTest._stop_web_server()
        MailpileSeleniumTest._stop_selenium_driver()

    def go_to_mailpile_home(self):
        self.driver.get('%s/in/inbox' % self._get_mailpile_url())

    def take_screenshot(self, filename):
        try:
            self.driver.save_screenshot(filename)  # save a screenshot to disk
        except WebDriverException:
            pass

    def dump_source_to(self, filename):
        with open(filename, 'w') as out:
            out.write(self.driver.page_source.encode('utf8'))

    def navigate_to(self, name):
        contacts = self.find_element_by_xpath(
            '//a[@alt="%s"]/span' % name)
        self.assertTrue(contacts.is_displayed())
        contacts.click()

    def submit_form(self, form_id):
        form = self.driver.find_element_by_id(form_id)
        form.submit()

    def fill_form_field(self, field, text):
        input_field = self.driver.find_element_by_name(field)
        input_field.send_keys(text)

    def assert_link_with_text(self, text):
        try:
            self.driver.find_element_by_link_text(text)
        except NoSuchElementException:
            raise AssertionError

    def click_element_with_link_text(self, text):
        try:
            self.driver.find_element_by_link_text(text).click()
        except NoSuchElementException:
            raise AssertionError

    def click_element_with_id(self, element_id):
        self.driver.find_element_by_id(element_id).click()

    def click_element_with_class(self, class_name):
        self.driver.find_element_by_class_name(class_name).click()

    def page_title(self):
        return self.driver.title

    def find_element_by_id(self, id):
        return self.driver.find_element_by_id(id)

    def find_element_containing_text(self, text):
        return self.driver.find_element_by_xpath("//*[contains(.,'%s')]" % text)

    def find_element_by_xpath(self, xpath):
        return self.driver.find_element_by_xpath(xpath)

    def find_element_by_class_name(self, class_name):
        return self.driver.find_element_by_class_name(class_name)

    def assert_text(self, text):
        self.find_element_containing_text(text)

    def wait_until_element_is_visible(self, element_id):
        self.wait_until_element_is_visible_by_locator((By.ID, element_id))

    def wait_until_element_is_visible_by_locator(self, locator_tuple):
        wait = WebDriverWait(self.driver, 10)
        wait.until(EC.visibility_of_element_located(locator_tuple))

    def wait_until_element_is_invisible_by_locator(self, locator_tuple):
        wait = WebDriverWait(self.driver, 10)
        wait.until(EC.invisibility_of_element_located(locator_tuple))

    def wait_until_element_has_class(self, locator_tuple, class_name):
        self.wait_for_element_condition(ElementHasClass(locator_tuple, class_name))

    def wait_until_element_has_not_class(self, locator_tuple, class_name):
        self.wait_for_element_condition(ElementHasNotClass(locator_tuple, class_name))

    def wait_for_element_condition(self, expected_conditions):
        wait = WebDriverWait(self.driver, 10)
        wait.until(expected_conditions)
<|endoftext|>


<commit_before>import argparse

from django.core.management.base import BaseCommand, CommandError

from ideascube.serveradmin.catalog import Catalog

class Command(BaseCommand):
    help = 'Manage apps and content'

    def add_arguments(self, parser):
        subs = parser.add_subparsers(
            title='Commands', dest='cmd', metavar='',
            parser_class=argparse.ArgumentParser)

        # -- Manage local cache -----------------------------------------------
        cache = subs.add_parser('cache', help='Manage cache')

        cachesubs = cache.add_subparsers(title='Commands', dest='cachecmd')
        cachesubs.required = True

        update = cachesubs.add_parser('update', help='Update the local cache')
        update.set_defaults(func=self.update_cache)

        clear = cachesubs.add_parser('clear', help='Clear the local cache')
        clear.set_defaults(func=self.clear_cache)

        # -- Manage remote sources --------------------------------------------
        remote = subs.add_parser('remotes', help='Manage remote sources')

        remotesubs = remote.add_subparsers(title='Commands', dest='remotecmd')
        remotesubs.required = True

        list = remotesubs.add_parser('list', help='List known remote sources')
        list.set_defaults(func=self.list_remotes)

        add = remotesubs.add_parser('add', help='Add a remote source')
        add.add_argument('id', help='The id for the new remote source')
        add.add_argument('name', help='The name for the new remote source')
        add.add_argument('url', help='The url for the new remote source')
        add.set_defaults(func=self.add_remote)

        rm = remotesubs.add_parser('remove', help='Remove a remote source')
        rm.add_argument('id', help='The id of the remote source to remove')
        rm.set_defaults(func=self.remove_remote)

        self.parser = parser

    def handle(self, *args, **options):
        if 'func' not in options:
            self.parser.print_help()
            self.parser.exit(1)

        self.catalog = Catalog()
        options['func'](options)

    # -- Manage local cache ---------------------------------------------------
    def update_cache(self, options):
        self.catalog.update_cache()

    def clear_cache(self, options):
        self.catalog.clear_cache()

    # -- Manage remote sources ------------------------------------------------
    def list_remotes(self, options):
        for remote in self.catalog.list_remotes():
            left = '[{0.id}] {0.name}'.format(remote)
            print('{0:>35} : {1.url}'.format(left, remote))

    def add_remote(self, options):
        try:
            self.catalog.add_remote(
                options['id'], options['name'], options['url'])

        except ValueError as e:
            raise CommandError(e)

        self.catalog.update_cache()

    def remove_remote(self, options):
        try:
            self.catalog.remove_remote(options['id'])

        except ValueError as e:
            raise CommandError(e)

        self.catalog.update_cache()
<commit_msg>catalog: Add package management to the cli<commit_after>import argparse

from django.core.management.base import BaseCommand, CommandError

from ideascube.serveradmin.catalog import Catalog, NoSuchPackage

class Command(BaseCommand):
    help = 'Manage apps and content'

    def add_arguments(self, parser):
        subs = parser.add_subparsers(
            title='Commands', dest='cmd', metavar='',
            parser_class=argparse.ArgumentParser)

        required_ids = argparse.ArgumentParser('required_ids', add_help=False)
        required_ids.add_argument(
            'ids', nargs='+', metavar='ID',
            help='One or more package ids (or globs)')

        optional_ids = argparse.ArgumentParser('optional_ids', add_help=False)
        optional_ids.add_argument(
            'ids', nargs='*', metavar='ID', default=['*'],
            help='Optional package ids (or globs)')

        # -- Manage content ---------------------------------------------------
        list = subs.add_parser(
            'list', parents=[optional_ids], help='List packages')
        group = list.add_mutually_exclusive_group()
        group.add_argument(
            '--installed', action='store_const', dest='filter',
            const='installed', help='Only list installed packages')
        group.add_argument(
            '--available', action='store_const', dest='filter',
            const='available', help='Only list available packages')
        group.add_argument(
            '--upgradable', action='store_const', dest='filter',
            const='upgradable',
            help='Only list installed packages with updates available')
        group.add_argument(
            '--all', action='store_const', dest='filter', const='all',
            help='List available and installed packages (default)')
        list.set_defaults(filter='all', func=self.list_packages)

        install = subs.add_parser(
            'install', parents=[required_ids],
            help='Install packages')
        install.set_defaults(func=self.install_packages)

        remove = subs.add_parser(
            'remove', parents=[required_ids], help='Remove packages')
        remove.set_defaults(func=self.remove_packages)

        upgrade = subs.add_parser(
            'upgrade', aliases=['update'],
            parents=[optional_ids], help='Upgrade packages')
        upgrade.set_defaults(func=self.upgrade_packages)

        # -- Manage local cache -----------------------------------------------
        cache = subs.add_parser('cache', help='Manage cache')

        cachesubs = cache.add_subparsers(title='Commands', dest='cachecmd')
        cachesubs.required = True

        update = cachesubs.add_parser('update', help='Update the local cache')
        update.set_defaults(func=self.update_cache)

        clear = cachesubs.add_parser('clear', help='Clear the local cache')
        clear.set_defaults(func=self.clear_cache)

        # -- Manage remote sources --------------------------------------------
        remote = subs.add_parser('remotes', help='Manage remote sources')

        remotesubs = remote.add_subparsers(title='Commands', dest='remotecmd')
        remotesubs.required = True

        list = remotesubs.add_parser('list', help='List known remote sources')
        list.set_defaults(func=self.list_remotes)

        add = remotesubs.add_parser('add', help='Add a remote source')
        add.add_argument('id', help='The id for the new remote source')
        add.add_argument('name', help='The name for the new remote source')
        add.add_argument('url', help='The url for the new remote source')
        add.set_defaults(func=self.add_remote)

        rm = remotesubs.add_parser('remove', help='Remove a remote source')
        rm.add_argument('id', help='The id of the remote source to remove')
        rm.set_defaults(func=self.remove_remote)

        self.parser = parser

    def handle(self, *args, **options):
        if 'func' not in options:
            self.parser.print_help()
            self.parser.exit(1)

        self.catalog = Catalog()
        options['func'](options)

    # -- Manage content -------------------------------------------------------
    def list_packages(self, options):
        fmt = ' {0.id:20}  {0.version:10}  {0.size:5}  {0.name}'

        if options['filter'] in ('all', 'installed'):
            pkgs = self.catalog.list_installed(options['ids'])

            if pkgs:
                print('Installed packages')

                for pkg in pkgs:
                    print(fmt.format(pkg))

        if options['filter'] in ('all', 'available'):
            pkgs = self.catalog.list_available(options['ids'])

            if pkgs:
                print('Available packages')

                for pkg in pkgs:
                    print(fmt.format(pkg))

        if options['filter'] == 'upgradable':
            pkgs = self.catalog.list_upgradable(options['ids'])

            if pkgs:
                print('Available updates')

                for pkg in pkgs:
                    print(fmt.format(pkg))

    def install_packages(self, options):
        try:
            self.catalog.install_packages(options['ids'])

        except NoSuchPackage as e:
            raise CommandError('No such package: {}'.format(e))

    def remove_packages(self, options):
        try:
            self.catalog.remove_packages(options['ids'])

        except NoSuchPackage as e:
            raise CommandError('No such package: {}'.format(e))

    def upgrade_packages(self, options):
        try:
            self.catalog.upgrade_packages(options['ids'])

        except NoSuchPackage as e:
            raise CommandError('No such package: {}'.format(e))

    # -- Manage local cache ---------------------------------------------------
    def update_cache(self, options):
        self.catalog.update_cache()

    def clear_cache(self, options):
        self.catalog.clear_cache()

    # -- Manage remote sources ------------------------------------------------
    def list_remotes(self, options):
        for remote in self.catalog.list_remotes():
            left = '[{0.id}] {0.name}'.format(remote)
            print('{0:>35} : {1.url}'.format(left, remote))

    def add_remote(self, options):
        try:
            self.catalog.add_remote(
                options['id'], options['name'], options['url'])

        except ValueError as e:
            raise CommandError(e)

        self.catalog.update_cache()

    def remove_remote(self, options):
        try:
            self.catalog.remove_remote(options['id'])

        except ValueError as e:
            raise CommandError(e)

        self.catalog.update_cache()
<|endoftext|>


<commit_before>"""
The :mod:`tslearn.metrics` module delivers time-series specific metrics to be
used at the core of machine learning algorithms.

**User guide:** See the :ref:`Dynamic Time Warping (DTW) <dtw>` section for
further details.
"""
from .dtw_variants import (dtw, dtw_limited_warping_length,
                           dtw_path_limited_warping_length, subsequence_path,
                           subsequence_cost_matrix,
                           dtw_path, dtw_path_from_metric,
                           dtw_subsequence_path, cdist_dtw,
                           GLOBAL_CONSTRAINT_CODE,
                           lb_envelope, lb_keogh)
from .sax import cdist_sax
from .softdtw_variants import (cdist_soft_dtw, cdist_gak,
                               cdist_soft_dtw_normalized, gak, soft_dtw,
                               sigma_gak, gamma_soft_dtw, SquaredEuclidean,
                               SoftDTW)

__author__ = 'Romain Tavenard romain.tavenard[at]univ-rennes2.fr'

TSLEARN_VALID_METRICS = ["dtw", "gak", "softdtw", "sax"]
VARIABLE_LENGTH_METRICS = ["dtw", "gak", "softdtw", "sax"]

__all__ = [
    TSLEARN_VALID_METRICS, VARIABLE_LENGTH_METRICS,
    dtw, dtw_limited_warping_length,
    dtw_path_limited_warping_length, subsequence_path,
    subsequence_cost_matrix,
    dtw_path, dtw_path_from_metric,
    dtw_subsequence_path, cdist_dtw,
    GLOBAL_CONSTRAINT_CODE,
    lb_envelope, lb_keogh,
    cdist_sax,
    cdist_soft_dtw, cdist_gak,
    cdist_soft_dtw_normalized, gak, soft_dtw,
    sigma_gak, gamma_soft_dtw, SquaredEuclidean, SoftDTW
]
<commit_msg>Fix missing import<commit_after>"""
The :mod:`tslearn.metrics` module delivers time-series specific metrics to be
used at the core of machine learning algorithms.

**User guide:** See the :ref:`Dynamic Time Warping (DTW) <dtw>` section for
further details.
"""
from .dtw_variants import (dtw, dtw_limited_warping_length,
                           dtw_path_limited_warping_length, subsequence_path,
                           subsequence_cost_matrix,
                           dtw_path, dtw_path_from_metric,
                           dtw_subsequence_path, cdist_dtw,
                           GLOBAL_CONSTRAINT_CODE,
                           lb_envelope, lb_keogh,
                           sakoe_chiba_mask, itakura_mask)
from .sax import cdist_sax
from .softdtw_variants import (cdist_soft_dtw, cdist_gak,
                               cdist_soft_dtw_normalized, gak, soft_dtw,
                               sigma_gak, gamma_soft_dtw, SquaredEuclidean,
                               SoftDTW)

__author__ = 'Romain Tavenard romain.tavenard[at]univ-rennes2.fr'

TSLEARN_VALID_METRICS = ["dtw", "gak", "softdtw", "sax"]
VARIABLE_LENGTH_METRICS = ["dtw", "gak", "softdtw", "sax"]

__all__ = [
    TSLEARN_VALID_METRICS, VARIABLE_LENGTH_METRICS,
    dtw, dtw_limited_warping_length,
    dtw_path_limited_warping_length, subsequence_path,
    subsequence_cost_matrix,
    dtw_path, dtw_path_from_metric,
    dtw_subsequence_path, cdist_dtw,
    GLOBAL_CONSTRAINT_CODE,
    lb_envelope, lb_keogh,
    sakoe_chiba_mask, itakura_mask,
    cdist_sax,
    cdist_soft_dtw, cdist_gak,
    cdist_soft_dtw_normalized, gak, soft_dtw,
    sigma_gak, gamma_soft_dtw, SquaredEuclidean, SoftDTW
]
<|endoftext|>