<commit_before>import unittest

from http_request_translator import translator

class TestTranslator(unittest.TestCase):

    ###
    # translator.parse_raw_request
    ###
    def test_parse_raw_request_http_version_with_path(self):
        for i in range(0, 10):
            for j in range(0, 10):
                raw_request = "GET /robots.txt HTTP/%d.%d\n"\
                              "Host: foo.bar" % (i, j)
                self.assertEqual(
                    translator.parse_raw_request(raw_request),
                    (
                        ['Host: foo.bar'],
                        {
                            'protocol': 'HTTP',
                            'pre_scheme': '',
                            'Host': 'foo.bar',
                            'version': '%d.%d' % (i, j),
                            'path': '/robots.txt',
                            'method': 'GET',
                            'data': ''
                        }
                    ),
                    'Invalid parsing of HTTP/%d.%d request!' % (i, j))

    def test_parse_raw_request_http_version_invalid(self):
        raw_request = "GET /\n"\
                      "Host: foo.bar"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: foo.bar'],
                {
                    'protocol': '',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '',
                    'path': '',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request!')
        raw_request = "GET / HTTP//1.b\n"\
                      "Host: foo.bar"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '/1.b',
                    'path': '/',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP//1.b request!')

    def test_parse_raw_request_http_with_parameter(self):
        raw_request = "GET /?foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/?foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with parameter in path!')

    def test_parse_raw_request_http_with_comment(self):
        raw_request = "GET /#foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/#foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with comment in path!')

    def test_parse_raw_request_http_with_coma(self):
        raw_request = "GET /;foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/;foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with coma in path!')

    def test_parse_raw_request_https_domain_no_port(self):
        raw_request = "GET https://google.com/robots.txt HTTP/1.1\n"\
                      "Host: google.com"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: google.com'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'google.com',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with domain host!')

    def test_parse_raw_request_https_domain_port(self):
        raw_request = "GET https://google.com:31337/robots.txt HTTP/1.1\n"\
                      "Host: google.com:31337"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: google.com:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'google.com:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with domain host and custom port!')

    def test_parse_raw_request_https_ipv4_no_port(self):
        raw_request = "GET https://127.0.0.1/robots.txt HTTP/1.1\n"\
                      "Host: 127.0.0.1"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: 127.0.0.1'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '127.0.0.1',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv4 host!')

    def test_parse_raw_request_https_ipv4_port(self):
        raw_request = "GET https://127.0.0.1:31337/robots.txt HTTP/1.1\n"\
                      "Host: 127.0.0.1:31337"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: 127.0.0.1:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '127.0.0.1:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv4 host and custom port!')

    def test_parse_raw_request_https_ipv6_no_port(self):
        raw_request = "GET https://[::1]/robots.txt HTTP/1.1\n"\
                      "Host: [::1]"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: [::1]'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '[::1]',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv6 host!')

    def test_parse_raw_request_https_ipv6_port(self):
        raw_request = "GET https://[::1]:31337/robots.txt HTTP/1.1\n"\
                      "Host: [::1]:31337"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                ['Host: [::1]:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '[::1]:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv6 host and custom port!')

    def test_parse_raw_request_multiple_host_header(self):
        raw_request = "GET https://foo.bar HTTP/1.1\n"\
                      "Host: foo.bar\n"\
                      "HoSt: foo.bar\n"\
                      "HOST: foo.bar \n"\
                      "host: foo.bar\n"\
                      "host:     foo.bar\n"
        self.assertEqual(
            translator.parse_raw_request(raw_request),
            (
                [
                    'Host: foo.bar',
                    'HoSt: foo.bar',
                    'HOST: foo.bar ',
                    'host: foo.bar',
                    'host:     foo.bar',
                ],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with multiple Host headers!')

    def test_parse_raw_request_invalid_header(self):
        raw_request = "GET https://foo.bar HTTP/1.1\n"\
                      "Host"
        with self.assertRaises(ValueError):
            translator.parse_raw_request(raw_request)

    def test_parse_raw_request_no_path(self):
        raw_request = "GET\n"\
                      "Host: foo.bar"
        with self.assertRaises(ValueError):
            translator.parse_raw_request(raw_request)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Add tests for _parse_request method<commit_after>import unittest

from http_request_translator.interface import HttpRequestTranslator

class TestTranslator(unittest.TestCase):

    ###
    # translator.parse_raw_request
    ###
    def test_parse_raw_request_http_version_with_path(self):
        for i in range(0, 10):
            for j in range(0, 10):
                raw_request = "GET /robots.txt HTTP/%d.%d\n"\
                              "Host: foo.bar" % (i, j)
                self.assertEqual(
                    HttpRequestTranslator(request=raw_request)._parse_request(),
                    (
                        ['Host: foo.bar'],
                        {
                            'protocol': 'HTTP',
                            'pre_scheme': '',
                            'Host': 'foo.bar',
                            'version': '%d.%d' % (i, j),
                            'path': '/robots.txt',
                            'method': 'GET',
                            'data': ''
                        }
                    ),
                    'Invalid parsing of HTTP/%d.%d request!' % (i, j))

    def test_parse_raw_request_http_version_invalid(self):
        raw_request = "GET /\n"\
                      "Host: foo.bar"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: foo.bar'],
                {
                    'protocol': '',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '',
                    'path': '',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request!')
        raw_request = "GET / HTTP//1.b\n"\
                      "Host: foo.bar"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '/1.b',
                    'path': '/',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP//1.b request!')

    def test_parse_raw_request_http_with_parameter(self):
        raw_request = "GET /?foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/?foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with parameter in path!')

    def test_parse_raw_request_http_with_comment(self):
        raw_request = "GET /#foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/#foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with comment in path!')

    def test_parse_raw_request_http_with_coma(self):
        raw_request = "GET /;foo=bar HTTP/1.1\n"\
                      "Host: foo.bar"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: foo.bar'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': '',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '/;foo=bar',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of request with coma in path!')

    def test_parse_raw_request_https_domain_no_port(self):
        raw_request = "GET https://google.com/robots.txt HTTP/1.1\n"\
                      "Host: google.com"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: google.com'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'google.com',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with domain host!')

    def test_parse_raw_request_https_domain_port(self):
        raw_request = "GET https://google.com:31337/robots.txt HTTP/1.1\n"\
                      "Host: google.com:31337"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: google.com:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'google.com:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with domain host and custom port!')

    def test_parse_raw_request_https_ipv4_no_port(self):
        raw_request = "GET https://127.0.0.1/robots.txt HTTP/1.1\n"\
                      "Host: 127.0.0.1"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: 127.0.0.1'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '127.0.0.1',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv4 host!')

    def test_parse_raw_request_https_ipv4_port(self):
        raw_request = "GET https://127.0.0.1:31337/robots.txt HTTP/1.1\n"\
                      "Host: 127.0.0.1:31337"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: 127.0.0.1:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '127.0.0.1:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv4 host and custom port!')

    def test_parse_raw_request_https_ipv6_no_port(self):
        raw_request = "GET https://[::1]/robots.txt HTTP/1.1\n"\
                      "Host: [::1]"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: [::1]'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '[::1]',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv6 host!')

    def test_parse_raw_request_https_ipv6_port(self):
        raw_request = "GET https://[::1]:31337/robots.txt HTTP/1.1\n"\
                      "Host: [::1]:31337"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                ['Host: [::1]:31337'],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': '[::1]:31337',
                    'version': '1.1',
                    'path': '/robots.txt',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with IPv6 host and custom port!')

    def test_parse_raw_request_multiple_host_header(self):
        raw_request = "GET https://foo.bar HTTP/1.1\n"\
                      "Host: foo.bar\n"\
                      "HoSt: foo.bar\n"\
                      "HOST: foo.bar \n"\
                      "host: foo.bar\n"\
                      "host:     foo.bar\n"
        self.assertEqual(
            HttpRequestTranslator(request=raw_request)._parse_request(),
            (
                [
                    'Host: foo.bar',
                    'HoSt: foo.bar',
                    'HOST: foo.bar ',
                    'host: foo.bar',
                    'host:     foo.bar',
                ],
                {
                    'protocol': 'HTTP',
                    'pre_scheme': 'https://',
                    'Host': 'foo.bar',
                    'version': '1.1',
                    'path': '',
                    'method': 'GET',
                    'data': ''
                }
            ),
            'Invalid parsing of HTTP request with multiple Host headers!')

    def test_parse_raw_request_invalid_header(self):
        raw_request = "GET https://foo.bar HTTP/1.1\n"\
                      "Host"
        with self.assertRaises(ValueError):
            HttpRequestTranslator(request=raw_request)._parse_request()

    def test_parse_raw_request_no_path(self):
        raw_request = "GET\n"\
                      "Host: foo.bar"
        with self.assertRaises(ValueError):
            HttpRequestTranslator(request=raw_request)._parse_request()

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""
byceps.blueprints.newsletter.views
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:Copyright: 2006-2016 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

from collections import namedtuple

from flask import abort, g, request

from ...util.framework import create_blueprint, flash_success
from ...util.templating import templated
from ...util.views import redirect_to

from .forms import SubscriptionForm
from .models import SubscriptionState
from . import service

blueprint = create_blueprint('newsletter', __name__)

@blueprint.route('/subscriptions/mine/update')
@templated
def subscription_update_form():
    """Show a form to update the current user's subscription."""
    user = get_current_user_or_404()

    state = service.get_subscription_state(user, g.party.brand)

    obj = namedtuple('Obj', 'state')(state.name)
    form = SubscriptionForm(obj=obj)

    return {
        'form': form,
    }

@blueprint.route('/subscriptions/mine', methods=['POST'])
def subscription_update():
    """Update the current user's subscription."""
    user = get_current_user_or_404()

    form = SubscriptionForm(request.form)

    state = SubscriptionState[form.state.data]

    if state is SubscriptionState.requested:
        service.subscribe(user, g.party.brand)
    elif state is SubscriptionState.declined:
        service.unsubscribe(user, g.party.brand)
    else:
        abort(400, 'Unknown subscription state.')

    flash_success('Dein Newsletter-Abonnement wurde aktualisiert.')
    return redirect_to('user.view_current')

def get_current_user_or_404():
    user = g.current_user
    if not user.is_active:
        abort(404)

    return user
<commit_msg>Display appropriate flash message on newsletter subscription state change based on the new state.<commit_after># -*- coding: utf-8 -*-

"""
byceps.blueprints.newsletter.views
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:Copyright: 2006-2016 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

from collections import namedtuple

from flask import abort, g, request

from ...util.framework import create_blueprint, flash_success
from ...util.templating import templated
from ...util.views import redirect_to

from .forms import SubscriptionForm
from .models import SubscriptionState
from . import service

blueprint = create_blueprint('newsletter', __name__)

@blueprint.route('/subscriptions/mine/update')
@templated
def subscription_update_form():
    """Show a form to update the current user's subscription."""
    user = get_current_user_or_404()

    state = service.get_subscription_state(user, g.party.brand)

    obj = namedtuple('Obj', 'state')(state.name)
    form = SubscriptionForm(obj=obj)

    return {
        'form': form,
    }

@blueprint.route('/subscriptions/mine', methods=['POST'])
def subscription_update():
    """Update the current user's subscription."""
    user = get_current_user_or_404()

    form = SubscriptionForm(request.form)

    state = SubscriptionState[form.state.data]

    if state is SubscriptionState.requested:
        service.subscribe(user, g.party.brand)
        flash_success('Du hast dich zum Newsletter angemeldet.')
    elif state is SubscriptionState.declined:
        service.unsubscribe(user, g.party.brand)
        flash_success('Du hast dich vom Newsletter abgemeldet.')
    else:
        abort(400, 'Unknown subscription state.')

    return redirect_to('user.view_current')

def get_current_user_or_404():
    user = g.current_user
    if not user.is_active:
        abort(404)

    return user
<|endoftext|>


<commit_before># Copyright 2020 Planet Labs, Inc.
# Copyright 2022 Planet Labs PBC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging

import pytest

from planet import specs

LOGGER = logging.getLogger(__name__)

TEST_PRODUCT_BUNDLE = 'visual'
ALL_PRODUCT_BUNDLES = [
    'analytic',
    'analytic_udm2',
    'analytic_3b_udm2',
    'analytic_5b',
    'analytic_5b_udm2',
    'analytic_8b_udm2',
    'visual',
    'uncalibrated_dn',
    'uncalibrated_dn_udm2',
    'basic_analytic',
    'basic_analytic_udm2',
    'basic_analytic_8b_udm2',
    'basic_uncalibrated_dn',
    'basic_uncalibrated_dn_udm2',
    'analytic_sr',
    'analytic_sr_udm2',
    'analytic_8b_sr_udm2',
    'basic_uncalibrated_dn_nitf',
    'basic_uncalibrated_dn_nitf_udm2',
    'basic_analytic_nitf',
    'basic_analytic_nitf_udm2',
    'basic_panchromatic',
    'basic_panchromatic_dn',
    'panchromatic',
    'panchromatic_dn',
    'panchromatic_dn_udm2',
    'pansharpened',
    'pansharpened_udm2',
    'basic_l1a_dn'
]
# must be a valid item type for TEST_PRODUCT_BUNDLE
TEST_ITEM_TYPE = 'PSScene'
ALL_ITEM_TYPES = [
    'PSOrthoTile',
    'Sentinel1',
    'REOrthoTile',
    'PSScene',
    'PSScene4Band',
    'Landsat8L1G',
    'PSScene3Band',
    'REScene',
    'MOD09GA',
    'MYD09GA',
    'MOD09GQ',
    'SkySatCollect',
    'Sentinel2L1C',
    'MYD09GQ',
    'SkySatScene'
]

def test_get_type_match():
    spec_list = ['Locket', 'drop', 'DEER']

    test_entry = 'locket'
    field_name = 'field_name'
    assert 'Locket' == specs.get_match(test_entry, spec_list, field_name)

    with pytest.raises(specs.SpecificationException):
        specs.get_match('a', ['b'], field_name)

def test_validate_bundle_supported():
    assert 'visual' == specs.validate_bundle(TEST_ITEM_TYPE, 'VISUAL')

def test_validate_bundle_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_bundle(TEST_ITEM_TYPE, 'notsupported')

def test_validate_bundle_notsupported_item_type():
    with pytest.raises(specs.SpecificationException):
        specs.validate_item_type('wha')

def test_validate_item_type_supported():
    assert 'PSOrthoTile' == specs.validate_item_type('psorthotile')

def test_validate_item_type_notsupported_itemtype():
    with pytest.raises(specs.SpecificationException):
        specs.validate_item_type('notsupported')

def test_validate_order_type_supported():
    assert 'full' == specs.validate_order_type('FULL')

def test_validate_order_type_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_order_type('notsupported')

def test_validate_arhive_type_supported():
    assert 'zip' == specs.validate_archive_type('ZIP')

def test_validate_arhive_type_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_archive_type('notsupported')

def test_validate_file_format_supported():
    assert 'COG' == specs.validate_file_format('cog')

def test_validate_file_format_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_archive_type('notsupported')

def test_get_product_bundles():
    bundles = specs.get_product_bundles()
    assert TEST_PRODUCT_BUNDLE in bundles

def test_get_item_types_with_bundle():
    item_types = specs.get_item_types(product_bundle=TEST_PRODUCT_BUNDLE)
    assert TEST_ITEM_TYPE in item_types

def test_get_item_types_without_bundle():
    item_types = specs.get_item_types()
    for item in item_types:
        assert item in ALL_ITEM_TYPES

def test_validate_supported_bundles_success():
    validated_bundle = specs.validate_supported_bundles(
        TEST_ITEM_TYPE, TEST_PRODUCT_BUNDLE, ALL_PRODUCT_BUNDLES)
    assert validated_bundle in ALL_PRODUCT_BUNDLES

def test_validate_supported_bundles_fail():
    with pytest.raises(specs.SpecificationException):
        specs.validate_supported_bundles(TEST_ITEM_TYPE,
                                         'analytic',
                                         ALL_PRODUCT_BUNDLES)
<commit_msg>Added a new test for product bundle with item type.<commit_after># Copyright 2020 Planet Labs, Inc.
# Copyright 2022 Planet Labs PBC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging

import pytest

from planet import specs

LOGGER = logging.getLogger(__name__)

TEST_PRODUCT_BUNDLE = 'visual'
ALL_PRODUCT_BUNDLES = [
    'analytic',
    'analytic_udm2',
    'analytic_3b_udm2',
    'analytic_5b',
    'analytic_5b_udm2',
    'analytic_8b_udm2',
    'visual',
    'uncalibrated_dn',
    'uncalibrated_dn_udm2',
    'basic_analytic',
    'basic_analytic_udm2',
    'basic_analytic_8b_udm2',
    'basic_uncalibrated_dn',
    'basic_uncalibrated_dn_udm2',
    'analytic_sr',
    'analytic_sr_udm2',
    'analytic_8b_sr_udm2',
    'basic_uncalibrated_dn_nitf',
    'basic_uncalibrated_dn_nitf_udm2',
    'basic_analytic_nitf',
    'basic_analytic_nitf_udm2',
    'basic_panchromatic',
    'basic_panchromatic_dn',
    'panchromatic',
    'panchromatic_dn',
    'panchromatic_dn_udm2',
    'pansharpened',
    'pansharpened_udm2',
    'basic_l1a_dn'
]
# must be a valid item type for TEST_PRODUCT_BUNDLE
TEST_ITEM_TYPE = 'PSScene'
ALL_ITEM_TYPES = [
    'PSOrthoTile',
    'Sentinel1',
    'REOrthoTile',
    'PSScene',
    'PSScene4Band',
    'Landsat8L1G',
    'PSScene3Band',
    'REScene',
    'MOD09GA',
    'MYD09GA',
    'MOD09GQ',
    'SkySatCollect',
    'Sentinel2L1C',
    'MYD09GQ',
    'SkySatScene'
]

def test_get_type_match():
    spec_list = ['Locket', 'drop', 'DEER']

    test_entry = 'locket'
    field_name = 'field_name'
    assert 'Locket' == specs.get_match(test_entry, spec_list, field_name)

    with pytest.raises(specs.SpecificationException):
        specs.get_match('a', ['b'], field_name)

def test_validate_bundle_supported():
    assert 'visual' == specs.validate_bundle(TEST_ITEM_TYPE, 'VISUAL')

def test_validate_bundle_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_bundle(TEST_ITEM_TYPE, 'notsupported')

def test_validate_bundle_notsupported_item_type():
    with pytest.raises(specs.SpecificationException):
        specs.validate_item_type('wha')

def test_validate_item_type_supported():
    assert 'PSOrthoTile' == specs.validate_item_type('psorthotile')

def test_validate_item_type_notsupported_itemtype():
    with pytest.raises(specs.SpecificationException):
        specs.validate_item_type('notsupported')

def test_validate_order_type_supported():
    assert 'full' == specs.validate_order_type('FULL')

def test_validate_order_type_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_order_type('notsupported')

def test_validate_arhive_type_supported():
    assert 'zip' == specs.validate_archive_type('ZIP')

def test_validate_arhive_type_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_archive_type('notsupported')

def test_validate_file_format_supported():
    assert 'COG' == specs.validate_file_format('cog')

def test_validate_file_format_notsupported():
    with pytest.raises(specs.SpecificationException):
        specs.validate_archive_type('notsupported')

def test_get_product_bundles_with_item_type():
    bundles = specs.get_product_bundles(item_type=TEST_ITEM_TYPE)
    assert TEST_PRODUCT_BUNDLE in bundles

def test_get_product_bundles_without_item_type():
    bundles = specs.get_product_bundles()
    for bundle in bundles:
        assert bundle in ALL_PRODUCT_BUNDLES

def test_get_item_types_with_bundle():
    item_types = specs.get_item_types(product_bundle=TEST_PRODUCT_BUNDLE)
    assert TEST_ITEM_TYPE in item_types

def test_get_item_types_without_bundle():
    item_types = specs.get_item_types()
    for item in item_types:
        assert item in ALL_ITEM_TYPES

def test_validate_supported_bundles_success():
    validated_bundle = specs.validate_supported_bundles(
        TEST_ITEM_TYPE, TEST_PRODUCT_BUNDLE, ALL_PRODUCT_BUNDLES)
    assert validated_bundle in ALL_PRODUCT_BUNDLES

def test_validate_supported_bundles_fail():
    with pytest.raises(specs.SpecificationException):
        specs.validate_supported_bundles(TEST_ITEM_TYPE,
                                         'analytic',
                                         ALL_PRODUCT_BUNDLES)
<|endoftext|>


<commit_before>import _path_pyblosxom

import string
import os
import os.path
from Pyblosxom import tools

class Testis_year:
    """tools.is_year"""

    def test_must_be_four_digits(self):
        assert tools.is_year("abab") == 0
        assert tools.is_year("ab") == 0
        assert tools.is_year("199") == 0
        assert tools.is_year("19999") == 0
        assert tools.is_year("1997") == 1
        assert tools.is_year("2097") == 1

    def test_must_start_with_19_or_20(self):
        assert tools.is_year("3090") == 0
        assert tools.is_year("0101") == 0

    def test_everything_else_returns_false(self):
        assert tools.is_year(None) == 0
        assert tools.is_year("") == 0
        assert tools.is_year("ab") == 0
        assert tools.is_year("97") == 0

class Testparse_args:
    """tools.parse_args"""
    def test_args_can_be_len_0(self):
        assert tools.parse_args( [] ) == []

    def test_single_args(self):
        assert tools.parse_args( ["--test"] ) == [ ("--test", "") ]
        assert tools.parse_args( ["-a", "-b", "--c"] ) == [ ("-a", ""),
                                                            ("-b", ""),
                                                            ("--c", "") ]
    def test_key_value_args(self):
        assert tools.parse_args( ["--test", "foo"] ) == [ ("--test", "foo") ]
        assert tools.parse_args( ["--test", "foo",
                                  "-a",
                                  "--bar", "baz"] ) == [ ("--test", "foo"),
                                                         ("-a", ""),
                                                         ("--bar", "baz") ]

class TestgenerateRandStr():
    """tools.generateRandStr

    Note: This is a mediocre test because generateRandStr produces
    a string that's of random length and random content.
    It's possible for this test to pass even when the code is bad.
    """

    def _gen_checker(self, s, minlen, maxlen):
        assert len(s) >= minlen and len(s) <= maxlen
        for c in s:
            assert c in string.letters or c in string.digits

    def test_generates_a_random_string(self):
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)

    def test_generates_a_random_string_between_minlen_and_maxlen(self):
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)

class TestEscapeText():
    """tools.escape_text"""
    def test_single_quote_to_pos(self):
        assert tools.escape_text("a'b") == "a'b"

    def test_double_quote_to_quot(self):
        assert tools.escape_text("a\"b") == "a"b"

    def test_everything_else_unchanged(self):
        assert tools.escape_text(None) == None
        assert tools.escape_text("") == ""
        assert tools.escape_text("abc") == "abc"

class TestUrlencodeText():
    """tools.urlencode_text"""
    def test_equals_to_3D(self):
        assert tools.urlencode_text("a=c") == "a%3Dc"

    def test_ampersand_to_26(self):
        assert tools.urlencode_text("a&c") == "a%26c"

    def test_space_to_20(self):
        assert tools.urlencode_text("a c") == "a%20c"

    def test_everything_else_unchanged(self):
        assert tools.urlencode_text(None) == None
        assert tools.urlencode_text("") == ""
        assert tools.urlencode_text("abc") == "abc"

class TestVariableDict():
    """tools.VariableDict class"""

    def __init__(self):
        self.vd = tools.VariableDict()
        self.vd.update( { "a": "b", \
                          "title": "This is my story", \
                          "title2": "Story: \"aha!\"", \
                          "filename": "joe.txt" } )

    def _compare_unordered_lists(self, lista, listb):
        assert len(lista) == len(listb)

        lista.sort()
        listb.sort()

        for mem in lista:
            print "comparing '" + mem + "' and '" + listb[0] + "'"
            assert mem == listb[0]
            listb.pop(0)

        assert len(listb) == 0

    def test_extends_dicts(self):
        assert self.vd.get("a") == "b"
        assert self.vd.get("title") == "This is my story"
        assert self.vd.get("title2") == "Story: \"aha!\""
        assert self.vd.get("badkey", "novalue") == "novalue"

        self._compare_unordered_lists(self.vd.keys(),
                                      ["a", "title", "title2", "filename"])
        self._compare_unordered_lists(self.vd.values(),
                                      ["b", "This is my story",
                                       "Story: \"aha!\"", "joe.txt"])

    def test_can_escape_value(self):
        assert self.vd.get("a_escaped") == "b"
        assert self.vd.get("title_escaped") == "This is my story"
        assert self.vd.get("title2_escaped") == "Story: "aha!""
        assert self.vd.get("badkey_escaped", "novalue") == "novalue"
        assert self.vd.get("badkey_escaped", "no value") == "no value"

    def test_can_urlencode_value(self):
        assert self.vd.get("a_urlencoded") == "b"
        assert self.vd.get("title_urlencoded") == "This%20is%20my%20story"
        assert self.vd.get("title2_urlencoded") == "Story%3A%20%22aha%21%22"
        assert self.vd.get("badkey_urlencoded", "novalue") == "novalue"
        assert self.vd.get("badkey_urlencoded", "no value") == "no%20value"

class TestStripper:
    """tools.Stripper class"""

    def _strip(self, text):
        s = tools.Stripper()
        s.feed(text)
        s.close()
        return s.gettext()

    def test_replaces_html_markup_from_string_with_space(self):
        s = tools.Stripper()
        assert self._strip("") == ""
        assert self._strip("abc") == "abc"
        assert self._strip("abc") == " abc "
        assert self._strip("abc def ghi") == "abc  def  ghi"

class TestWhatExt:
    """tools.what_ext"""
    def __init__(self):
        self._files = ["a.txt", "b.html", "c.txtl"]
        d = os.path.dirname(__file__)
        self._d = os.path.join(d, "what_ext_test_dir")

    def _setup(self):
        """
        Creates the directory with some files in it.
        """
        try:
            os.mkdir(self._d)

            for mem in self._files:
                f = open(os.path.join(self._d, mem), "w")
                f.write("lorem ipsum")
                f.close()
        except:
            self._teardown()

    def _teardown(self):
        """
        Cleans up the test files and the directory that we created.
        """
        for mem in self._files:
            try:
                os.remove(os.path.join(self._d, mem))
            except:
                pass

        try:
            os.rmdir(self._d)
        except:
            pass

    def test_returns_extension_if_file_has_extension(self):
        self._setup()
        try:
            assert "txt" == tools.what_ext(["txt", "html"],
                                           os.path.join(self._d, "a"))
            assert "html" == tools.what_ext(["txt", "html"],
                                           os.path.join(self._d, "b"))
        finally:
            self._teardown()

    def test_returns_None_if_extension_not_present(self):
        self._setup()
        try:
            assert None == tools.what_ext([],
                                          os.path.join(self._d, "a"))
            assert None == tools.what_ext(["html"],
                                          os.path.join(self._d, "a"))
        finally:
            self._teardown()
<commit_msg>Added additional tests.  Added some minor documentation.  Fixed a test.<commit_after>import _path_pyblosxom

import string
import os
import os.path
from Pyblosxom import tools

class Testis_year:
    """tools.is_year"""

    def test_must_be_four_digits(self):
        assert tools.is_year("abab") == 0
        assert tools.is_year("ab") == 0
        assert tools.is_year("199") == 0
        assert tools.is_year("19999") == 0
        assert tools.is_year("1997") == 1
        assert tools.is_year("2097") == 1

    def test_must_start_with_19_or_20(self):
        assert tools.is_year("3090") == 0
        assert tools.is_year("0101") == 0

    def test_everything_else_returns_false(self):
        assert tools.is_year(None) == 0
        assert tools.is_year("") == 0
        assert tools.is_year("ab") == 0
        assert tools.is_year("97") == 0

class Testparse_args:
    """tools.parse_args"""
    def test_args_can_be_len_0(self):
        assert tools.parse_args( [] ) == []

    def test_single_args(self):
        assert tools.parse_args( ["--test"] ) == [ ("--test", "") ]
        assert tools.parse_args( ["-a", "-b", "--c"] ) == [ ("-a", ""),
                                                            ("-b", ""),
                                                            ("--c", "") ]
    def test_key_value_args(self):
        assert tools.parse_args( ["--test", "foo"] ) == [ ("--test", "foo") ]
        assert tools.parse_args( ["--test", "foo",
                                  "-a",
                                  "--bar", "baz"] ) == [ ("--test", "foo"),
                                                         ("-a", ""),
                                                         ("--bar", "baz") ]

class TestgenerateRandStr():
    """tools.generateRandStr

    Note: This is a mediocre test because generateRandStr produces
    a string that's of random length and random content.
    It's possible for this test to pass even when the code is bad.
    """

    def _gen_checker(self, s, minlen, maxlen):
        assert len(s) >= minlen and len(s) <= maxlen
        for c in s:
            assert c in string.letters or c in string.digits

    def test_generates_a_random_string(self):
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)
        self._gen_checker(tools.generateRandStr(), 5, 10)

    def test_generates_a_random_string_between_minlen_and_maxlen(self):
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(4, 10), 4, 10)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)
        self._gen_checker(tools.generateRandStr(3, 12), 3, 12)

class TestEscapeText():
    """tools.escape_text"""
    def test_none_to_none(self):
        assert tools.escape_text(None) == None

    def test_empty_string_to_empty_string(self):
        assert tools.escape_text("") == ""

    def test_single_quote_to_pos(self):
        assert tools.escape_text("a'b") == "a'b"

    def test_double_quote_to_quot(self):
        assert tools.escape_text("a\"b") == "a"b"

    def test_everything_else_unchanged(self):
        assert tools.escape_text(None) == None
        assert tools.escape_text("") == ""
        assert tools.escape_text("abc") == "abc"

class TestUrlencodeText():
    """tools.urlencode_text"""
    def test_none_to_none(self):
        assert tools.urlencode_text(None) == None

    def test_empty_string_to_empty_string(self):
        assert tools.urlencode_text("") == ""

    def test_equals_to_3D(self):
        assert tools.urlencode_text("a=c") == "a%3Dc"

    def test_ampersand_to_26(self):
        assert tools.urlencode_text("a&c") == "a%26c"

    def test_space_to_20(self):
        assert tools.urlencode_text("a c") == "a%20c"

    def test_everything_else_unchanged(self):
        assert tools.urlencode_text(None) == None
        assert tools.urlencode_text("") == ""
        assert tools.urlencode_text("abc") == "abc"

class TestVariableDict():
    """tools.VariableDict class"""

    def __init__(self):
        # we use the same VariableDict instance through all
        # the tests since none of them mutate the instance.
        self.vd = tools.VariableDict()
        self.vd.update( { "a": "b", \
                          "title": "This is my story", \
                          "title2": "Story: \"aha!\"", \
                          "filename": "joe.txt" } )

    def _compare_unordered_lists(self, lista, listb):
        assert len(lista) == len(listb)

        lista.sort()
        listb.sort()

        for mem in lista:
            print "comparing '" + mem + "' and '" + listb[0] + "'"
            assert mem == listb[0]
            listb.pop(0)

        assert len(listb) == 0

    def test_extends_dicts(self):
        assert self.vd.get("a") == "b"
        assert self.vd.get("title") == "This is my story"
        assert self.vd.get("title2") == "Story: \"aha!\""
        assert self.vd.get("badkey", "novalue") == "novalue"

        self._compare_unordered_lists(self.vd.keys(),
                                      ["a", "title", "title2", "filename"])
        self._compare_unordered_lists(self.vd.values(),
                                      ["b", "This is my story",
                                       "Story: \"aha!\"", "joe.txt"])

    def test_can_escape_value(self):
        assert self.vd.get("a_escaped") == "b"
        assert self.vd.get("title_escaped") == "This is my story"
        assert self.vd.get("title2_escaped") == "Story: "aha!""
        assert self.vd.get("badkey_escaped", "novalue") == "novalue"
        assert self.vd.get("badkey_escaped", 'no"value') == 'no"value'

    def test_can_urlencode_value(self):
        assert self.vd.get("a_urlencoded") == "b"
        assert self.vd.get("title_urlencoded") == "This%20is%20my%20story"
        assert self.vd.get("title2_urlencoded") == "Story%3A%20%22aha%21%22"
        assert self.vd.get("badkey_urlencoded", "novalue") == "novalue"
        assert self.vd.get("badkey_urlencoded", "no value") == "no value"

class TestStripper:
    """tools.Stripper class"""

    def _strip(self, text):
        s = tools.Stripper()
        s.feed(text)
        s.close()
        return s.gettext()

    def test_replaces_html_markup_from_string_with_space(self):
        s = tools.Stripper()
        assert self._strip("") == ""
        assert self._strip("abc") == "abc"
        assert self._strip("abc") == " abc "
        assert self._strip("abc
") == "abc "
        assert self._strip("abc def ghi") == "abc  def  ghi"

class TestWhatExt:
    """tools.what_ext"""
    def __init__(self):
        self._files = ["a.txt", "b.html", "c.txtl"]
        d = os.path.dirname(__file__)
        self._d = os.path.join(d, "what_ext_test_dir")

    def _setup(self):
        """
        Creates the directory with some files in it.
        """
        try:
            os.mkdir(self._d)

            for mem in self._files:
                f = open(os.path.join(self._d, mem), "w")
                f.write("lorem ipsum")
                f.close()
        except:
            self._teardown()

    def _teardown(self):
        """
        Cleans up the test files and the directory that we created.
        """
        for mem in self._files:
            try:
                os.remove(os.path.join(self._d, mem))
            except:
                pass

        try:
            os.rmdir(self._d)
        except:
            pass

    def test_returns_extension_if_file_has_extension(self):
        self._setup()
        try:
            assert "txt" == tools.what_ext(["txt", "html"],
                                           os.path.join(self._d, "a"))
            assert "html" == tools.what_ext(["txt", "html"],
                                           os.path.join(self._d, "b"))
        finally:
            self._teardown()

    def test_returns_None_if_extension_not_present(self):
        self._setup()
        try:
            assert None == tools.what_ext([],
                                          os.path.join(self._d, "a"))
            assert None == tools.what_ext(["html"],
                                          os.path.join(self._d, "a"))
        finally:
            self._teardown()
<|endoftext|>


<commit_before>from wagtail_personalisation.utils import (
    exclude_variants, impersonate_other_page)

from wagtail.core.models import Page as WagtailPage

class Page(object):
    def __init__(self, path, depth, url_path, title):
        self.path = path
        self.depth = depth
        self.url_path = url_path
        self.title = title

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

def test_impersonate_other_page():
    page = Page(path="/", depth=0, url_path="/", title="Hoi")
    other_page = Page(path="/other", depth=1, url_path="/other", title="Doei")

    impersonate_other_page(page, other_page)

    assert page == other_page

class Metadata(object):
    def __init__(self, is_canonical=True):
        self.is_canonical = is_canonical

class PersonalisationMetadataPage(object):
    def __init__(self):
        self.personalisation_metadata = Metadata()

def test_exclude_variants_includes_pages_with_no_metadata_property():
    page = PersonalisationMetadataPage()
    del page.personalisation_metadata
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_includes_pages_with_metadata_none():
    page = PersonalisationMetadataPage()
    page.personalisation_metadata = None
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_includes_pages_with_metadata_canonical():
    page = PersonalisationMetadataPage()
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_excludes_pages_with_metadata_not_canonical():
    page = PersonalisationMetadataPage()
    page.personalisation_metadata.is_canonical = False
    result = exclude_variants([page])
    assert result == []

def test_exclude_variants_with_pages_querysets():
    '''
    Test that excludes variant works for querysets too
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result == pages

def test_exclude_variants_with_pages_querysets_not_canonical():
    '''
    Test that excludes variant works for querysets too
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')
    # add variants
    for page in pages:
        page.personalisation_metadata = Metadata(is_canonical=False)
        page.save()

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result.count() == 0

def test_exclude_variants_with_pages_querysets_meta_none():
    '''
    Test that excludes variant works for querysets too
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')
    # add variants
    for page in pages:
        page.personalisation_metadata = None
        page.save()

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result == pages
<commit_msg>fixed some comments<commit_after>from wagtail_personalisation.utils import (
    exclude_variants, impersonate_other_page)

from wagtail.core.models import Page as WagtailPage

class Page(object):
    def __init__(self, path, depth, url_path, title):
        self.path = path
        self.depth = depth
        self.url_path = url_path
        self.title = title

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

def test_impersonate_other_page():
    page = Page(path="/", depth=0, url_path="/", title="Hoi")
    other_page = Page(path="/other", depth=1, url_path="/other", title="Doei")

    impersonate_other_page(page, other_page)

    assert page == other_page

class Metadata(object):
    def __init__(self, is_canonical=True):
        self.is_canonical = is_canonical

class PersonalisationMetadataPage(object):
    def __init__(self):
        self.personalisation_metadata = Metadata()

def test_exclude_variants_includes_pages_with_no_metadata_property():
    page = PersonalisationMetadataPage()
    del page.personalisation_metadata
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_includes_pages_with_metadata_none():
    page = PersonalisationMetadataPage()
    page.personalisation_metadata = None
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_includes_pages_with_metadata_canonical():
    page = PersonalisationMetadataPage()
    result = exclude_variants([page])
    assert result == [page]

def test_exclude_variants_excludes_pages_with_metadata_not_canonical():
    page = PersonalisationMetadataPage()
    page.personalisation_metadata.is_canonical = False
    result = exclude_variants([page])
    assert result == []

def test_exclude_variants_with_pages_querysets():
    '''
    Test that excludes variant works for querysets
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result == pages

def test_exclude_variants_with_pages_querysets_not_canonical():
    '''
    Test that excludes variant works for querysets with
    personalisation_metadata canonical False
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')
    # add variants
    for page in pages:
        page.personalisation_metadata = Metadata(is_canonical=False)
        page.save()

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result.count() == 0

def test_exclude_variants_with_pages_querysets_meta_none():
    '''
    Test that excludes variant works for querysets with meta as none
    '''
    for i in range(5):
        page = WagtailPage(path="/" + str(i), depth=0, url_path="/", title="Hoi " + str(i))
        page.save()
    pages = WagtailPage.objects.all().order_by('id')
    # add variants
    for page in pages:
        page.personalisation_metadata = None
        page.save()

    result = exclude_variants(pages)
    assert type(result) == type(pages)
    assert result == pages
<|endoftext|>


<commit_before>from datetime import datetime
from github3.utils import stream_response_to_file, timestamp_parameter

import io
import mock
import pytest
import requests

class TestTimestampConverter:
    def test_datetimes(self):
        timestamp = datetime(2010, 6, 1, 12, 15, 30)
        assert '2010-06-01T12:15:30' == timestamp_parameter(timestamp)

    def test_valid_datestring(self):
        testvals = (
            '2010-06-01',
            '2010-06-01T12:15:30',
            '2010-06-01T12:14:30.12321+02:00',
            '2010-06-01T12:14:30.12321-02:00',
            '2010-06-01T12:14:30.2115Z',
        )
        for timestamp in testvals:
            assert timestamp == timestamp_parameter(timestamp)

    def test_invalid_datestring(self):
        testvals = (
            '2012-16-04',
            '2012-06-01v!',
            'fish',
            '2010-06-01T12:14:30.12321+02',
            '2010-06-01T12:70:30.12321+02',
        )
        for timestamp in testvals:
            pytest.raises(ValueError, timestamp_parameter, timestamp)

    def test_none_handling(self):
        assert timestamp_parameter(None, allow_none=True) is None
        pytest.raises(ValueError, timestamp_parameter, None,
                      allow_none=False)

    def test_invalid_type_handling(self):
        pytest.raises(ValueError, timestamp_parameter, 1)

@pytest.fixture
def mocked_open():
    return mock.mock_open()

@pytest.fixture
def response():
    r = requests.Response()
    r.raw = io.BytesIO(b'fake data')
    r.headers.update({'content-disposition': 'filename=a_file_name'})
    return r

class OpenFile:
    def __init__(self):
        self.data = b''
        self.written_to = False

    def write(self, data):
        self.written_to = True
        self.data += data

class TestStreamingDownloads:
    def test_opens_a_new_file(self, mocked_open, response):
        with mock.patch('github3.utils.open', mocked_open, create=True):
            stream_response_to_file(response, 'some_file')

        mocked_open.assert_called_once_with('some_file', 'wb')

    def test_uses_existing_file(self, response):
        fd = OpenFile()
        stream_response_to_file(response, fd)
        assert fd.written_to is True
        assert fd.data == b'fake data'

    def test_finds_filename_in_headers(self, mocked_open, response):
        with mock.patch('github3.utils.open', mocked_open, create=True):
            stream_response_to_file(response)

        mocked_open.assert_called_once_with('a_file_name', 'wb')
<commit_msg>Add more assertions around the mocked file<commit_after>from datetime import datetime
from github3.utils import stream_response_to_file, timestamp_parameter

import io
import mock
import pytest
import requests

class TestTimestampConverter:
    def test_datetimes(self):
        timestamp = datetime(2010, 6, 1, 12, 15, 30)
        assert '2010-06-01T12:15:30' == timestamp_parameter(timestamp)

    def test_valid_datestring(self):
        testvals = (
            '2010-06-01',
            '2010-06-01T12:15:30',
            '2010-06-01T12:14:30.12321+02:00',
            '2010-06-01T12:14:30.12321-02:00',
            '2010-06-01T12:14:30.2115Z',
        )
        for timestamp in testvals:
            assert timestamp == timestamp_parameter(timestamp)

    def test_invalid_datestring(self):
        testvals = (
            '2012-16-04',
            '2012-06-01v!',
            'fish',
            '2010-06-01T12:14:30.12321+02',
            '2010-06-01T12:70:30.12321+02',
        )
        for timestamp in testvals:
            pytest.raises(ValueError, timestamp_parameter, timestamp)

    def test_none_handling(self):
        assert timestamp_parameter(None, allow_none=True) is None
        pytest.raises(ValueError, timestamp_parameter, None,
                      allow_none=False)

    def test_invalid_type_handling(self):
        pytest.raises(ValueError, timestamp_parameter, 1)

@pytest.fixture
def mocked_open():
    return mock.mock_open()

@pytest.fixture
def response():
    r = requests.Response()
    r.raw = io.BytesIO(b'fake data')
    r.headers.update({'content-disposition': 'filename=a_file_name'})
    return r

class OpenFile:
    def __init__(self):
        self.data = b''
        self.written_to = False

    def write(self, data):
        self.written_to = True
        self.data += data

class TestStreamingDownloads:
    def test_opens_a_new_file(self, mocked_open, response):
        with mock.patch('github3.utils.open', mocked_open, create=True):
            stream_response_to_file(response, 'some_file')

        mocked_open.assert_called_once_with('some_file', 'wb')
        mocked_open().write.assert_called_once_with(b'fake data')
        mocked_open().close.assert_called_once_with()

    def test_uses_existing_file(self, response):
        fd = OpenFile()
        stream_response_to_file(response, fd)
        assert fd.written_to is True
        assert fd.data == b'fake data'

    def test_finds_filename_in_headers(self, mocked_open, response):
        with mock.patch('github3.utils.open', mocked_open, create=True):
            stream_response_to_file(response)

        mocked_open.assert_called_once_with('a_file_name', 'wb')
        mocked_open().write.assert_called_once_with(b'fake data')
        mocked_open().close.assert_called_once_with()
<|endoftext|>


<commit_before>"""
A HTTP fetcher class.
"""

import os
import urlparse
import requests
import shutil
import mimetypes
from rfc6266 import parse_requests_response

from tea_party.fetchers.base_fetcher import BaseFetcher

class HttpFetcher(BaseFetcher):

    """
    Fetchs a file from the web, using HTTP or HTTPS.
    """

    shortname = 'http'

    def normalize_location(self, location):
        """
        Checks that the `location` is an HTTP(S) URL.
        """

        url = urlparse.urlparse(location)

        if url.scheme in ['http', 'https']:
            return urlparse.urlunparse(url)

    def do_fetch(self, target):
        """
        Fetch the archive at the specified location.
        """

        response = requests.get(self.location, stream=True)
        response.raise_for_status()

        mimetype = response.headers.get('content-type')
        encoding = response.headers.get('content-encoding')

        extension = mimetypes.guess_extension(mimetype)
        content_disposition = parse_requests_response(response)

        filename = content_disposition.filename_sanitized(extension=extension[1:], default_filename='archive')

        content_length = response.headers.get('content-length')

        if content_length is not None:
            content_length = int(content_length)

        target_file_path = os.path.join(target, filename)

        self.on_start(target=os.path.basename(target_file_path), size=content_length)

        with open(target_file_path, 'wb') as target_file:
            current_size = 0

            for buf in response.iter_content(1024):

                if buf:
                    target_file.write(buf)
                    current_size += len(buf)

                    self.on_update(progress=current_size)

        self.on_finish()

        return {
            'archive_path': target_file_path,
            'mimetype': mimetype,
            'encoding': encoding,
        }
<commit_msg>Fixed HTTP fetcher downloads when mimetype is not known<commit_after>"""
A HTTP fetcher class.
"""

import os
import urlparse
import requests
import shutil
import mimetypes
from rfc6266 import parse_requests_response

from tea_party.log import LOGGER
from tea_party.fetchers.base_fetcher import BaseFetcher

class HttpFetcher(BaseFetcher):

    """
    Fetchs a file from the web, using HTTP or HTTPS.
    """

    shortname = 'http'

    def normalize_location(self, location):
        """
        Checks that the `location` is an HTTP(S) URL.
        """

        url = urlparse.urlparse(location)

        if url.scheme in ['http', 'https']:
            return urlparse.urlunparse(url)

    def do_fetch(self, target):
        """
        Fetch the archive at the specified location.
        """

        response = requests.get(self.location, stream=True)
        response.raise_for_status()

        mimetype = response.headers.get('content-type')
        encoding = response.headers.get('content-encoding')

        extension = mimetypes.guess_extension(mimetype)

        if not extension:
            LOGGER.debug('No extension registered for this mimetype (%s). Guessing one from the URL...', mimetype)

            extension = os.path.splitext(urlparse.urlparse(self.location).path)[1]

        if extension and extension.startswith('.'):
            extension = extension[1:]

        content_disposition = parse_requests_response(response)

        filename = content_disposition.filename_sanitized(extension=extension, default_filename='archive')

        content_length = response.headers.get('content-length')

        if content_length is not None:
            content_length = int(content_length)

        target_file_path = os.path.join(target, filename)

        self.on_start(target=os.path.basename(target_file_path), size=content_length)

        with open(target_file_path, 'wb') as target_file:
            current_size = 0

            for buf in response.iter_content(1024):

                if buf:
                    target_file.write(buf)
                    current_size += len(buf)

                    self.on_update(progress=current_size)

        self.on_finish()

        return {
            'archive_path': target_file_path,
            'mimetype': mimetype,
            'encoding': encoding,
        }
<|endoftext|>


<commit_before>import asyncio
import itertools
import time

from .chatgetter import ChatGetter
from ... import utils, errors

class Conversation(ChatGetter):
    """
    Represents a conversation inside an specific chat.

    A conversation keeps track of new messages since it was
    created until its exit and easily lets you query the
    current state.

    If you need a conversation across two or more chats,
    you should use two conversations and synchronize them
    as you better see fit.
    """
    _id_counter = 0
    _custom_counter = 0

    def __init__(self, client, input_chat,
                 *, timeout, total_timeout, max_messages,
                 exclusive, replies_are_responses):
        self._id = Conversation._id_counter
        Conversation._id_counter += 1

        self._client = client
        self._chat = None
        self._input_chat = input_chat
        self._chat_peer = None
        self._broadcast = None

        self._timeout = timeout
        self._total_timeout = total_timeout
        self._total_due = None

        self._outgoing = set()
        self._last_outgoing = 0
        self._incoming = []
        self._last_incoming = 0
        self._max_incoming = max_messages
        self._last_read = None
        self._custom = {}

        self._pending_responses = {}
        self._pending_replies = {}
        self._pending_edits = {}
        self._pending_reads = {}

        self._exclusive = exclusive

        # The user is able to expect two responses for the same message.
        # {desired message ID: next incoming index}
        self._response_indices = {}
        if replies_are_responses:
            self._reply_indices = self._response_indices
        else:
            self._reply_indices = {}

        self._edit_dates = {}

    async def send_message(self, *args, **kwargs):
        """
        Sends a message in the context of this conversation. Shorthand
        for `telethon.client.messages.MessageMethods.send_message` with
        ``entity`` already set.
        """
        message = await self._client.send_message(
            self._input_chat, *args, **kwargs)

        self._outgoing.add(message.id)
        self._last_outgoing = message.id
        return message

    async def send_file(self, *args, **kwargs):
        """
        Sends a file in the context of this conversation. Shorthand
        for `telethon.client.uploads.UploadMethods.send_file` with
        ``entity`` already set.
        """
        message = await self._client.send_file(
            self._input_chat, *args, **kwargs)

        self._outgoing.add(message.id)
        self._last_outgoing = message.id
        return message

    def mark_read(self, message=None):
        """
        Marks as read the latest received message if ``message is None``.
        Otherwise, marks as read until the given message (or message ID).

        This is equivalent to calling `client.send_read_acknowledge
        <telethon.client.messages.MessageMethods.send_read_acknowledge>`.
        """
        if message is None:
            if self._incoming:
                message = self._incoming[-1].id
            else:
                message = 0
        elif not isinstance(message, int):
            message = message.id

        return self._client.send_read_acknowledge(
            self._input_chat, max_id=message)

    async def get_response(self, message=None, *, timeout=None):
        """
        Returns a coroutine that will resolve once a response arrives.

        Args:
            message (`Message <telethon.tl.custom.message.Message>` | `int`, optional):
                The message (or the message ID) for which a response
                is expected. By default this is the last sent message.

            timeout (`int` | `float`, optional):
                If present, this `timeout` will override the
                per-action timeout defined for the conversation.
        """
        return await self._get_message(
            message, self._response_indices, self._pending_responses, timeout,
            lambda x, y: True
        )

    async def get_reply(self, message=None, *, timeout=None):
        """
        Returns a coroutine that will resolve once a reply
        (that is, a message being a reply) arrives. The
        arguments are the same as those for `get_response`.
        """
        return await self._get_message(
            message, self._reply_indices, self._pending_replies, timeout,
            lambda x, y: x.reply_to_msg_id == y
        )

    def _get_message(
            self, target_message, indices, pending, timeout, condition):
        """
        Gets the next desired message under the desired condition.

        Args:
            target_message (`object`):
                The target message for which we want to find another
                response that applies based on `condition`.

            indices (`dict`):
                This dictionary remembers the last ID chosen for the
                input `target_message`.

            pending (`dict`):
                This dictionary remembers {msg_id: Future} to be set
                once `condition` is met.

            timeout (`int`):
                The timeout override to use for this operation.

            condition (`callable`):
                The condition callable that checks if an incoming
                message is a valid response.
        """
        start_time = time.time()
        target_id = self._get_message_id(target_message)

        # If there is no last-chosen ID, make sure to pick one *after*
        # the input message, since we don't want responses back in time
        if target_id not in indices:
            for i, incoming in enumerate(self._incoming):
                if incoming.id > target_id:
                    indices[target_id] = i
                    break
            else:
                indices[target_id] = len(self._incoming)

        # We will always return a future from here, even if the result
        # can be set immediately. Otherwise, needing to await only
        # sometimes is an annoying edge case (i.e. we would return
        # a `Message` but `get_response()` always `await`'s).
        future = self._client.loop.create_future()

        # If there are enough responses saved return the next one
        last_idx = indices[target_id]
        if last_idx < len(self._incoming):
            incoming = self._incoming[last_idx]
            if condition(incoming, target_id):
                indices[target_id] += 1
                future.set_result(incoming)
                return future

        # Otherwise the next incoming response will be the one to use
        pending[target_id] = future
        return self._get_result(future, start_time, timeout)

    async def get_edit(self, message=None, *, timeout=None):
        """
        Awaits for an edit after the last message to arrive.
        The arguments are the same as those for `get_response`.
        """
        start_time = time.time()
        target_id = self._get_message_id(message)

        target_date = self._edit_dates.get(target_id, 0)
        earliest_edit = min(
            (x for x in self._incoming
             if x.edit_date
             and x.id > target_id
             and x.edit_date.timestamp() > target_date
             ),
            key=lambda x: x.edit_date.timestamp(),
            default=None
        )

        if earliest_edit and earliest_edit.edit_date.timestamp() > target_date:
            self._edit_dates[target_id] = earliest_edit.edit_date.timestamp()
            return earliest_edit

        # Otherwise the next incoming response will be the one to use
        future = asyncio.Future(loop=self._client.loop)
        self._pending_edits[target_id] = future
        return await self._get_result(future, start_time, timeout)

    async def wait_read(self, message=None, *, timeout=None):
        """
        Awaits for the sent message to be read. Note that receiving
        a response doesn't imply the message was read, and this action
        will also trigger even without a response.
        """
        start_time = time.time()
        future = self._client.loop.create_future()
        target_id = self._get_message_id(message)

        if self._last_read is None:
            self._last_read = target_id - 1

        if self._last_read >= target_id:
            return

        self._pending_reads[target_id] = future
        return await self._get_result(future, start_time, timeout)

    async def wait_event(self, event, *, timeout=None):
        """
        Waits for a custom event to occur. Timeouts still apply.

        Unless you're certain that your code will run fast enough,
        generally you should get a "handle" of this special coroutine
        before acting. Generally, you should do this:

        >>> from telethon import TelegramClient, events
        >>>
        >>> client = TelegramClient(...)
        >>>
        >>> async def main():
        >>>     async with client.conversation(...) as conv:
        >>>         response = conv.wait_event(events.NewMessage(incoming=True))
        >>>         await conv.send_message('Hi')
        >>>         response = await response

        This way your event can be registered before acting,
        since the response may arrive before your event was
        registered. It depends on your use case since this
        also means the event can arrive before you send
        a previous action.
        """
        start_time = time.time()
        if isinstance(event, type):
            event = event()

        await event.resolve(self._client)

        counter = Conversation._custom_counter
        Conversation._custom_counter += 1

        future = asyncio.Future(loop=self._client.loop)

        # We need the `async def` here because we want to block on the future
        # from `_get_result` by using `await` on it. If we returned the future
        # immediately we would `del` from `_custom` too early.

        async def result():
            try:
                return await self._get_result(future, start_time, timeout)
            finally:
                del self._custom[counter]

        self._custom[counter] = (event, future)
        return await result()

    async def _check_custom(self, built):
        for i, (ev, fut) in self._custom.items():
            ev_type = type(ev)
            if built[ev_type] and ev.filter(built[ev_type]):
                fut.set_result(built[ev_type])

    def _on_new_message(self, response):
        response = response.message
        if response.chat_id != self.chat_id or response.out:
            return

        if len(self._incoming) == self._max_incoming:
            self._cancel_all(ValueError('Too many incoming messages'))
            return

        self._incoming.append(response)

        found = []
        for msg_id in self._pending_responses:
            found.append(msg_id)
            self._response_indices[msg_id] = len(self._incoming)

        for msg_id in found:
            self._pending_responses.pop(msg_id).set_result(response)

        found.clear()
        for msg_id in self._pending_replies:
            if msg_id == response.reply_to_msg_id:
                found.append(msg_id)
                self._reply_indices[msg_id] = len(self._incoming)

        for msg_id in found:
            self._pending_replies.pop(msg_id).set_result(response)

    def _on_edit(self, message):
        message = message.message
        if message.chat_id != self.chat_id or message.out:
            return

        found = []
        for msg_id, pending in self._pending_edits.items():
            if msg_id < message.id:
                found.append(msg_id)
                self._edit_dates[msg_id] = message.edit_date.timestamp()

        for msg_id in found:
            self._pending_edits.pop(msg_id).set_result(message)

    def _on_read(self, event):
        if event.chat_id != self.chat_id or event.inbox:
            return

        self._last_read = event.max_id

        remove_reads = []
        for msg_id, pending in self._pending_reads.items():
            if msg_id >= self._last_read:
                remove_reads.append(msg_id)
                pending.set_result(True)

        for to_remove in remove_reads:
            del self._pending_reads[to_remove]

    def _get_message_id(self, message):
        if message:
            return message if isinstance(message, int) else message.id
        elif self._last_outgoing:
            return self._last_outgoing
        else:
            raise ValueError('No message was sent previously')

    def _get_result(self, future, start_time, timeout):
        due = self._total_due
        if timeout is None:
            timeout = self._timeout

        if timeout is not None:
            due = min(due, start_time + timeout)

        return asyncio.wait_for(
            future,
            timeout=None if due == float('inf') else due - time.time(),
            loop=self._client.loop
        )

    def _cancel_all(self, exception=None):
        for pending in itertools.chain(
                self._pending_responses.values(),
                self._pending_replies.values(),
                self._pending_edits.values()):
            if exception:
                pending.set_exception(exception)
            else:
                pending.cancel()

        for _, fut in self._custom.values():
            if exception:
                fut.set_exception(exception)
            else:
                fut.cancel()

    def __enter__(self):
        if self._client.loop.is_running():
            raise RuntimeError(
                'You must use "async with" if the event loop '
                'is running (i.e. you are inside an "async def")'
            )

        return self._client.loop.run_until_complete(self.__aenter__())

    async def __aenter__(self):
        self._input_chat = \
            await self._client.get_input_entity(self._input_chat)

        self._chat_peer = utils.get_peer(self._input_chat)

        # Make sure we're the only conversation in this chat if it's exclusive
        chat_id = utils.get_peer_id(self._chat_peer)
        count = self._client._ids_in_conversations.get(chat_id, 0)
        if self._exclusive and count:
            raise errors.AlreadyInConversationError()

        self._client._ids_in_conversations[chat_id] = count + 1
        self._client._conversations[self._id] = self

        self._last_outgoing = 0
        self._last_incoming = 0
        for d in (
                self._outgoing, self._incoming,
                self._pending_responses, self._pending_replies,
                self._pending_edits, self._response_indices,
                self._reply_indices, self._edit_dates, self._custom):
            d.clear()

        if self._total_timeout:
            self._total_due = time.time() + self._total_timeout
        else:
            self._total_due = float('inf')

        return self

    def cancel(self):
        """Cancels the current conversation and exits the context manager."""
        raise _ConversationCancelled()

    def __exit__(self, *args):
        return self._client.loop.run_until_complete(self.__aexit__(*args))

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        chat_id = utils.get_peer_id(self._chat_peer)
        if self._client._ids_in_conversations[chat_id] == 1:
            del self._client._ids_in_conversations[chat_id]
        else:
            self._client._ids_in_conversations[chat_id] -= 1

        del self._client._conversations[self._id]
        self._cancel_all()
        return isinstance(exc_val, _ConversationCancelled)

class _ConversationCancelled(InterruptedError):
    pass
<commit_msg>Work around message edits arriving too early in conversations<commit_after>import asyncio
import itertools
import time

from .chatgetter import ChatGetter
from ... import utils, errors

# Sometimes the edits arrive very fast (within the same second).
# In that case we add a small delta so that the age is older, for
# comparision purposes. This value is enough for up to 1000 messages.
_EDIT_COLLISION_DELTA = 0.001

class Conversation(ChatGetter):
    """
    Represents a conversation inside an specific chat.

    A conversation keeps track of new messages since it was
    created until its exit and easily lets you query the
    current state.

    If you need a conversation across two or more chats,
    you should use two conversations and synchronize them
    as you better see fit.
    """
    _id_counter = 0
    _custom_counter = 0

    def __init__(self, client, input_chat,
                 *, timeout, total_timeout, max_messages,
                 exclusive, replies_are_responses):
        self._id = Conversation._id_counter
        Conversation._id_counter += 1

        self._client = client
        self._chat = None
        self._input_chat = input_chat
        self._chat_peer = None
        self._broadcast = None

        self._timeout = timeout
        self._total_timeout = total_timeout
        self._total_due = None

        self._outgoing = set()
        self._last_outgoing = 0
        self._incoming = []
        self._last_incoming = 0
        self._max_incoming = max_messages
        self._last_read = None
        self._custom = {}

        self._pending_responses = {}
        self._pending_replies = {}
        self._pending_edits = {}
        self._pending_reads = {}

        self._exclusive = exclusive

        # The user is able to expect two responses for the same message.
        # {desired message ID: next incoming index}
        self._response_indices = {}
        if replies_are_responses:
            self._reply_indices = self._response_indices
        else:
            self._reply_indices = {}

        self._edit_dates = {}

    async def send_message(self, *args, **kwargs):
        """
        Sends a message in the context of this conversation. Shorthand
        for `telethon.client.messages.MessageMethods.send_message` with
        ``entity`` already set.
        """
        message = await self._client.send_message(
            self._input_chat, *args, **kwargs)

        self._outgoing.add(message.id)
        self._last_outgoing = message.id
        return message

    async def send_file(self, *args, **kwargs):
        """
        Sends a file in the context of this conversation. Shorthand
        for `telethon.client.uploads.UploadMethods.send_file` with
        ``entity`` already set.
        """
        message = await self._client.send_file(
            self._input_chat, *args, **kwargs)

        self._outgoing.add(message.id)
        self._last_outgoing = message.id
        return message

    def mark_read(self, message=None):
        """
        Marks as read the latest received message if ``message is None``.
        Otherwise, marks as read until the given message (or message ID).

        This is equivalent to calling `client.send_read_acknowledge
        <telethon.client.messages.MessageMethods.send_read_acknowledge>`.
        """
        if message is None:
            if self._incoming:
                message = self._incoming[-1].id
            else:
                message = 0
        elif not isinstance(message, int):
            message = message.id

        return self._client.send_read_acknowledge(
            self._input_chat, max_id=message)

    async def get_response(self, message=None, *, timeout=None):
        """
        Returns a coroutine that will resolve once a response arrives.

        Args:
            message (`Message <telethon.tl.custom.message.Message>` | `int`, optional):
                The message (or the message ID) for which a response
                is expected. By default this is the last sent message.

            timeout (`int` | `float`, optional):
                If present, this `timeout` will override the
                per-action timeout defined for the conversation.
        """
        return await self._get_message(
            message, self._response_indices, self._pending_responses, timeout,
            lambda x, y: True
        )

    async def get_reply(self, message=None, *, timeout=None):
        """
        Returns a coroutine that will resolve once a reply
        (that is, a message being a reply) arrives. The
        arguments are the same as those for `get_response`.
        """
        return await self._get_message(
            message, self._reply_indices, self._pending_replies, timeout,
            lambda x, y: x.reply_to_msg_id == y
        )

    def _get_message(
            self, target_message, indices, pending, timeout, condition):
        """
        Gets the next desired message under the desired condition.

        Args:
            target_message (`object`):
                The target message for which we want to find another
                response that applies based on `condition`.

            indices (`dict`):
                This dictionary remembers the last ID chosen for the
                input `target_message`.

            pending (`dict`):
                This dictionary remembers {msg_id: Future} to be set
                once `condition` is met.

            timeout (`int`):
                The timeout override to use for this operation.

            condition (`callable`):
                The condition callable that checks if an incoming
                message is a valid response.
        """
        start_time = time.time()
        target_id = self._get_message_id(target_message)

        # If there is no last-chosen ID, make sure to pick one *after*
        # the input message, since we don't want responses back in time
        if target_id not in indices:
            for i, incoming in enumerate(self._incoming):
                if incoming.id > target_id:
                    indices[target_id] = i
                    break
            else:
                indices[target_id] = len(self._incoming)

        # We will always return a future from here, even if the result
        # can be set immediately. Otherwise, needing to await only
        # sometimes is an annoying edge case (i.e. we would return
        # a `Message` but `get_response()` always `await`'s).
        future = self._client.loop.create_future()

        # If there are enough responses saved return the next one
        last_idx = indices[target_id]
        if last_idx < len(self._incoming):
            incoming = self._incoming[last_idx]
            if condition(incoming, target_id):
                indices[target_id] += 1
                future.set_result(incoming)
                return future

        # Otherwise the next incoming response will be the one to use
        pending[target_id] = future
        return self._get_result(future, start_time, timeout)

    async def get_edit(self, message=None, *, timeout=None):
        """
        Awaits for an edit after the last message to arrive.
        The arguments are the same as those for `get_response`.
        """
        start_time = time.time()
        target_id = self._get_message_id(message)

        target_date = self._edit_dates.get(target_id, 0)
        earliest_edit = min(
            (x for x in self._incoming
             if x.edit_date
             and x.id > target_id
             and x.edit_date.timestamp() > target_date
             ),
            key=lambda x: x.edit_date.timestamp(),
            default=None
        )

        if earliest_edit and earliest_edit.edit_date.timestamp() > target_date:
            self._edit_dates[target_id] = earliest_edit.edit_date.timestamp()
            return earliest_edit

        # Otherwise the next incoming response will be the one to use
        future = asyncio.Future(loop=self._client.loop)
        self._pending_edits[target_id] = future
        return await self._get_result(future, start_time, timeout)

    async def wait_read(self, message=None, *, timeout=None):
        """
        Awaits for the sent message to be read. Note that receiving
        a response doesn't imply the message was read, and this action
        will also trigger even without a response.
        """
        start_time = time.time()
        future = self._client.loop.create_future()
        target_id = self._get_message_id(message)

        if self._last_read is None:
            self._last_read = target_id - 1

        if self._last_read >= target_id:
            return

        self._pending_reads[target_id] = future
        return await self._get_result(future, start_time, timeout)

    async def wait_event(self, event, *, timeout=None):
        """
        Waits for a custom event to occur. Timeouts still apply.

        Unless you're certain that your code will run fast enough,
        generally you should get a "handle" of this special coroutine
        before acting. Generally, you should do this:

        >>> from telethon import TelegramClient, events
        >>>
        >>> client = TelegramClient(...)
        >>>
        >>> async def main():
        >>>     async with client.conversation(...) as conv:
        >>>         response = conv.wait_event(events.NewMessage(incoming=True))
        >>>         await conv.send_message('Hi')
        >>>         response = await response

        This way your event can be registered before acting,
        since the response may arrive before your event was
        registered. It depends on your use case since this
        also means the event can arrive before you send
        a previous action.
        """
        start_time = time.time()
        if isinstance(event, type):
            event = event()

        await event.resolve(self._client)

        counter = Conversation._custom_counter
        Conversation._custom_counter += 1

        future = asyncio.Future(loop=self._client.loop)

        # We need the `async def` here because we want to block on the future
        # from `_get_result` by using `await` on it. If we returned the future
        # immediately we would `del` from `_custom` too early.

        async def result():
            try:
                return await self._get_result(future, start_time, timeout)
            finally:
                del self._custom[counter]

        self._custom[counter] = (event, future)
        return await result()

    async def _check_custom(self, built):
        for i, (ev, fut) in self._custom.items():
            ev_type = type(ev)
            if built[ev_type] and ev.filter(built[ev_type]):
                fut.set_result(built[ev_type])

    def _on_new_message(self, response):
        response = response.message
        if response.chat_id != self.chat_id or response.out:
            return

        if len(self._incoming) == self._max_incoming:
            self._cancel_all(ValueError('Too many incoming messages'))
            return

        self._incoming.append(response)

        found = []
        for msg_id in self._pending_responses:
            found.append(msg_id)
            self._response_indices[msg_id] = len(self._incoming)

        for msg_id in found:
            self._pending_responses.pop(msg_id).set_result(response)

        found.clear()
        for msg_id in self._pending_replies:
            if msg_id == response.reply_to_msg_id:
                found.append(msg_id)
                self._reply_indices[msg_id] = len(self._incoming)

        for msg_id in found:
            self._pending_replies.pop(msg_id).set_result(response)

    def _on_edit(self, message):
        message = message.message
        if message.chat_id != self.chat_id or message.out:
            return

        found = []
        for msg_id, pending in self._pending_edits.items():
            if msg_id < message.id:
                found.append(msg_id)
                edit_ts = message.edit_date.timestamp()

                # We compare <= because edit_ts resolution is always to
                # seconds, but we may have increased _edit_dates before.
                # Since the dates are ever growing this is not a problem.
                if edit_ts <= self._edit_dates.get(msg_id, 0):
                    self._edit_dates[msg_id] += _EDIT_COLLISION_DELTA
                else:
                    self._edit_dates[msg_id] = message.edit_date.timestamp()

        for msg_id in found:
            self._pending_edits.pop(msg_id).set_result(message)

    def _on_read(self, event):
        if event.chat_id != self.chat_id or event.inbox:
            return

        self._last_read = event.max_id

        remove_reads = []
        for msg_id, pending in self._pending_reads.items():
            if msg_id >= self._last_read:
                remove_reads.append(msg_id)
                pending.set_result(True)

        for to_remove in remove_reads:
            del self._pending_reads[to_remove]

    def _get_message_id(self, message):
        if message:
            return message if isinstance(message, int) else message.id
        elif self._last_outgoing:
            return self._last_outgoing
        else:
            raise ValueError('No message was sent previously')

    def _get_result(self, future, start_time, timeout):
        due = self._total_due
        if timeout is None:
            timeout = self._timeout

        if timeout is not None:
            due = min(due, start_time + timeout)

        return asyncio.wait_for(
            future,
            timeout=None if due == float('inf') else due - time.time(),
            loop=self._client.loop
        )

    def _cancel_all(self, exception=None):
        for pending in itertools.chain(
                self._pending_responses.values(),
                self._pending_replies.values(),
                self._pending_edits.values()):
            if exception:
                pending.set_exception(exception)
            else:
                pending.cancel()

        for _, fut in self._custom.values():
            if exception:
                fut.set_exception(exception)
            else:
                fut.cancel()

    def __enter__(self):
        if self._client.loop.is_running():
            raise RuntimeError(
                'You must use "async with" if the event loop '
                'is running (i.e. you are inside an "async def")'
            )

        return self._client.loop.run_until_complete(self.__aenter__())

    async def __aenter__(self):
        self._input_chat = \
            await self._client.get_input_entity(self._input_chat)

        self._chat_peer = utils.get_peer(self._input_chat)

        # Make sure we're the only conversation in this chat if it's exclusive
        chat_id = utils.get_peer_id(self._chat_peer)
        count = self._client._ids_in_conversations.get(chat_id, 0)
        if self._exclusive and count:
            raise errors.AlreadyInConversationError()

        self._client._ids_in_conversations[chat_id] = count + 1
        self._client._conversations[self._id] = self

        self._last_outgoing = 0
        self._last_incoming = 0
        for d in (
                self._outgoing, self._incoming,
                self._pending_responses, self._pending_replies,
                self._pending_edits, self._response_indices,
                self._reply_indices, self._edit_dates, self._custom):
            d.clear()

        if self._total_timeout:
            self._total_due = time.time() + self._total_timeout
        else:
            self._total_due = float('inf')

        return self

    def cancel(self):
        """Cancels the current conversation and exits the context manager."""
        raise _ConversationCancelled()

    def __exit__(self, *args):
        return self._client.loop.run_until_complete(self.__aexit__(*args))

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        chat_id = utils.get_peer_id(self._chat_peer)
        if self._client._ids_in_conversations[chat_id] == 1:
            del self._client._ids_in_conversations[chat_id]
        else:
            self._client._ids_in_conversations[chat_id] -= 1

        del self._client._conversations[self._id]
        self._cancel_all()
        return isinstance(exc_val, _ConversationCancelled)

class _ConversationCancelled(InterruptedError):
    pass
<|endoftext|>