<commit_before># -*- coding: utf-8 -*-
"""
    Test suite for the util module
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import re
import unittest

from pygments import util, console

class FakeLexer(object):
    def analyse(text):
        return float(text)
    analyse = util.make_analysator(analyse)

class UtilTest(unittest.TestCase):

    def test_getoptions(self):
        raises = self.assertRaises
        equals = self.assertEqual

        equals(util.get_bool_opt({}, 'a', True), True)
        equals(util.get_bool_opt({}, 'a', 1), True)
        equals(util.get_bool_opt({}, 'a', 'true'), True)
        equals(util.get_bool_opt({}, 'a', 'no'), False)
        raises(util.OptionError, util.get_bool_opt, {}, 'a', [])
        raises(util.OptionError, util.get_bool_opt, {}, 'a', 'foo')

        equals(util.get_int_opt({}, 'a', 1), 1)
        raises(util.OptionError, util.get_int_opt, {}, 'a', [])
        raises(util.OptionError, util.get_int_opt, {}, 'a', 'bar')

        equals(util.get_list_opt({}, 'a', [1]), [1])
        equals(util.get_list_opt({}, 'a', '1 2'), ['1', '2'])
        raises(util.OptionError, util.get_list_opt, {}, 'a', 1)

    def test_docstring_headline(self):
        def f1():
            """
            docstring headline

            other text
            """
        def f2():
            """
            docstring
            headline

            other text
            """

        self.assertEqual(util.docstring_headline(f1), "docstring headline")
        self.assertEqual(util.docstring_headline(f2), "docstring headline")

    def test_analysator_returns_float(self):
        # If an analysator wrapped by make_analysator returns a floating point
        # number, then that number will be returned by the wrapper.
        self.assertEqual(FakeLexer.analyse('0.5'), 0.5)

    def test_analysator_returns_boolean(self):
        # If an analysator wrapped by make_analysator returns a boolean value,
        # then the wrapper will return 1.0 if the boolean was True or 0.0 if
        # it was False.
        self.assertEqual(FakeLexer.analyse(True), 1.0)
        self.assertEqual(FakeLexer.analyse(False), 0.0)

    def test_analysator_raises_exception(self):
        # If an analysator wrapped by make_analysator raises an exception,
        # then the wrapper will return 0.0.
        class ErrorLexer(object):
            def analyse(text):
                raise RuntimeError('something bad happened')
            analyse = util.make_analysator(analyse)
        self.assertEqual(ErrorLexer.analyse(''), 0.0)

    def test_analysator_value_error(self):
        # When converting the analysator's return value to a float a
        # ValueError may occur.  If that happens 0.0 is returned instead.
        self.assertEqual(FakeLexer.analyse('bad input'), 0.0)

    def test_analysator_type_error(self):
        # When converting the analysator's return value to a float a
        # TypeError may occur.  If that happens 0.0 is returned instead.
        self.assertEqual(FakeLexer.analyse(None), 0.0)

    def test_shebang_matches(self):
        self.assertTrue(util.shebang_matches('#!/usr/bin/env python', r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!/usr/bin/python2.4', r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!/usr/bin/startsomethingwith python',
                                             r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!C:\\Python2.4\\Python.exe',
                                             r'python(2\.\d)?'))

        self.assertFalse(util.shebang_matches('#!/usr/bin/python-ruby',
                                              r'python(2\.\d)?'))
        self.assertFalse(util.shebang_matches('#!/usr/bin/python/ruby',
                                              r'python(2\.\d)?'))
        self.assertFalse(util.shebang_matches('#!', r'python'))

    def test_doctype_matches(self):
        self.assertTrue(util.doctype_matches(
            '<!DOCTYPE html PUBLIC "a"> ', 'html.*'))
        self.assertFalse(util.doctype_matches(
            '<?xml ?> <DOCTYPE html PUBLIC "a"> ', 'html.*'))
        self.assertTrue(util.html_doctype_matches(
            '<?xml ?><!DOCTYPE html PUBLIC  "-//W3C//DTD XHTML 1.0 Strict//EN">'))

    def test_xml(self):
        self.assertTrue(util.looks_like_xml(
            '<?xml ?><!DOCTYPE html PUBLIC  "-//W3C//DTD XHTML 1.0 Strict//EN">'))
        self.assertTrue(util.looks_like_xml('abc'))
        self.assertFalse(util.looks_like_xml(''))

    def test_unirange(self):
        first_non_bmp = u'\U00010000'
        r = re.compile(util.unirange(0x10000, 0x20000))
        m = r.match(first_non_bmp)
        self.assertTrue(m)
        self.assertEqual(m.end(), len(first_non_bmp))
        self.assertFalse(r.match(u'\uffff'))
        self.assertFalse(r.match(u'xxx'))
        # Tests that end is inclusive
        r = re.compile(util.unirange(0x10000, 0x10000) + '+')
        # Tests that the plus works for the entire unicode point, if narrow
        # build
        m = r.match(first_non_bmp * 2)
        self.assertTrue(m)
        self.assertEqual(m.end(), len(first_non_bmp) * 2)

    def test_format_lines(self):
        lst = ['cat', 'dog']
        output = util.format_lines('var', lst)
        d = {}
        exec(output, d)
        self.assertTrue(isinstance(d['var'], tuple))
        self.assertEqual(('cat', 'dog'), d['var'])

    def test_duplicates_removed_seq_types(self):
        # tuple
        x = util.duplicates_removed(('a', 'a', 'b'))
        self.assertEqual(['a', 'b'], x)
        # list
        x = util.duplicates_removed(['a', 'a', 'b'])
        self.assertEqual(['a', 'b'], x)
        # iterator
        x = util.duplicates_removed(iter(('a', 'a', 'b')))
        self.assertEqual(['a', 'b'], x)

    def test_duplicates_removed_nonconsecutive(self):
        # keeps first
        x = util.duplicates_removed(('a', 'b', 'a'))
        self.assertEqual(['a', 'b'], x)

class ConsoleTest(unittest.TestCase):

    def test_ansiformat(self):
        f = console.ansiformat
        c = console.codes
        all_attrs = f('+*_blue_*+', 'text')
        self.assertTrue(c['blue'] in all_attrs and c['blink'] in all_attrs
                        and c['bold'] in all_attrs and c['underline'] in all_attrs
                        and c['reset'] in all_attrs)
        self.assertRaises(KeyError, f, '*mauve*', 'text')

    def test_functions(self):
        self.assertEqual(console.reset_color(), console.codes['reset'])
        self.assertEqual(console.colorize('blue', 'text'),
                         console.codes['blue'] + 'text' + console.codes['reset'])
<commit_msg>Add more coverage for pygments.util.<commit_after># -*- coding: utf-8 -*-
"""
    Test suite for the util module
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import re
import unittest

from pygments import util, console

class FakeLexer(object):
    def analyse(text):
        return text
    analyse = util.make_analysator(analyse)

class UtilTest(unittest.TestCase):

    def test_getoptions(self):
        raises = self.assertRaises
        equals = self.assertEqual

        equals(util.get_bool_opt({}, 'a', True), True)
        equals(util.get_bool_opt({}, 'a', 1), True)
        equals(util.get_bool_opt({}, 'a', 'true'), True)
        equals(util.get_bool_opt({}, 'a', 'no'), False)
        raises(util.OptionError, util.get_bool_opt, {}, 'a', [])
        raises(util.OptionError, util.get_bool_opt, {}, 'a', 'foo')

        equals(util.get_int_opt({}, 'a', 1), 1)
        raises(util.OptionError, util.get_int_opt, {}, 'a', [])
        raises(util.OptionError, util.get_int_opt, {}, 'a', 'bar')

        equals(util.get_list_opt({}, 'a', [1]), [1])
        equals(util.get_list_opt({}, 'a', '1 2'), ['1', '2'])
        raises(util.OptionError, util.get_list_opt, {}, 'a', 1)

        equals(util.get_choice_opt({}, 'a', ['foo', 'bar'], 'bar'), 'bar')
        equals(util.get_choice_opt({}, 'a', ['foo', 'bar'], 'Bar', True), 'bar')
        raises(util.OptionError, util.get_choice_opt, {}, 'a',
               ['foo', 'bar'], 'baz')

    def test_docstring_headline(self):
        def f1():
            """
            docstring headline

            other text
            """
        def f2():
            """
            docstring
            headline

            other text
            """
        def f3():
            pass

        self.assertEqual(util.docstring_headline(f1), 'docstring headline')
        self.assertEqual(util.docstring_headline(f2), 'docstring headline')
        self.assertEqual(util.docstring_headline(f3), '')

    def test_analysator_returns_float(self):
        # If an analysator wrapped by make_analysator returns a floating point
        # number, then that number will be returned by the wrapper.
        self.assertEqual(FakeLexer.analyse('0.5'), 0.5)

    def test_analysator_returns_boolean(self):
        # If an analysator wrapped by make_analysator returns a boolean value,
        # then the wrapper will return 1.0 if the boolean was True or 0.0 if
        # it was False.
        self.assertEqual(FakeLexer.analyse(True), 1.0)
        self.assertEqual(FakeLexer.analyse(False), 0.0)

    def test_analysator_raises_exception(self):
        # If an analysator wrapped by make_analysator raises an exception,
        # then the wrapper will return 0.0.
        class ErrorLexer(object):
            def analyse(text):
                raise RuntimeError('something bad happened')
            analyse = util.make_analysator(analyse)
        self.assertEqual(ErrorLexer.analyse(''), 0.0)

    def test_analysator_value_error(self):
        # When converting the analysator's return value to a float a
        # ValueError may occur.  If that happens 0.0 is returned instead.
        self.assertEqual(FakeLexer.analyse('bad input'), 0.0)

    def test_analysator_type_error(self):
        # When converting the analysator's return value to a float a
        # TypeError may occur.  If that happens 0.0 is returned instead.
        self.assertEqual(FakeLexer.analyse('xxx'), 0.0)

    def test_shebang_matches(self):
        self.assertTrue(util.shebang_matches('#!/usr/bin/env python\n', r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!/usr/bin/python2.4', r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!/usr/bin/startsomethingwith python',
                                             r'python(2\.\d)?'))
        self.assertTrue(util.shebang_matches('#!C:\\Python2.4\\Python.exe',
                                             r'python(2\.\d)?'))

        self.assertFalse(util.shebang_matches('#!/usr/bin/python-ruby',
                                              r'python(2\.\d)?'))
        self.assertFalse(util.shebang_matches('#!/usr/bin/python/ruby',
                                              r'python(2\.\d)?'))
        self.assertFalse(util.shebang_matches('#!', r'python'))

    def test_doctype_matches(self):
        self.assertTrue(util.doctype_matches(
            '<!DOCTYPE html> ', 'html.*'))
        self.assertFalse(util.doctype_matches(
            '<?xml ?> <DOCTYPE html PUBLIC "a"> ', 'html.*'))
        self.assertTrue(util.html_doctype_matches(
            '<?xml ?><!DOCTYPE html PUBLIC  "-//W3C//DTD XHTML 1.0 Strict//EN">'))

    def test_xml(self):
        self.assertTrue(util.looks_like_xml(
            '<?xml ?><!DOCTYPE html PUBLIC  "-//W3C//DTD XHTML 1.0 Strict//EN">'))
        self.assertTrue(util.looks_like_xml('abc'))
        self.assertFalse(util.looks_like_xml(''))

    def test_unirange(self):
        first_non_bmp = u'\U00010000'
        r = re.compile(util.unirange(0x10000, 0x20000))
        m = r.match(first_non_bmp)
        self.assertTrue(m)
        self.assertEqual(m.end(), len(first_non_bmp))
        self.assertFalse(r.match(u'\uffff'))
        self.assertFalse(r.match(u'xxx'))
        # Tests that end is inclusive
        r = re.compile(util.unirange(0x10000, 0x10000) + '+')
        # Tests that the plus works for the entire unicode point, if narrow
        # build
        m = r.match(first_non_bmp * 2)
        self.assertTrue(m)
        self.assertEqual(m.end(), len(first_non_bmp) * 2)

    def test_format_lines(self):
        lst = ['cat', 'dog']
        output = util.format_lines('var', lst)
        d = {}
        exec(output, d)
        self.assertTrue(isinstance(d['var'], tuple))
        self.assertEqual(('cat', 'dog'), d['var'])

    def test_duplicates_removed_seq_types(self):
        # tuple
        x = util.duplicates_removed(('a', 'a', 'b'))
        self.assertEqual(['a', 'b'], x)
        # list
        x = util.duplicates_removed(['a', 'a', 'b'])
        self.assertEqual(['a', 'b'], x)
        # iterator
        x = util.duplicates_removed(iter(('a', 'a', 'b')))
        self.assertEqual(['a', 'b'], x)

    def test_duplicates_removed_nonconsecutive(self):
        # keeps first
        x = util.duplicates_removed(('a', 'b', 'a'))
        self.assertEqual(['a', 'b'], x)

    def test_guess_decode(self):
        # UTF-8 should be decoded as UTF-8
        s = util.guess_decode(u'\xff'.encode('utf-8'))
        self.assertEqual(s, (u'\xff', 'utf-8'))

        # otherwise, it could be latin1 or the locale encoding...
        import locale
        s = util.guess_decode(b'\xff')
        self.assertTrue(s[1] in ('latin1', locale.getpreferredencoding()))

    def test_guess_decode_from_terminal(self):
        class Term:
            encoding = 'utf-7'

        s = util.guess_decode_from_terminal(u'\xff'.encode('utf-7'), Term)
        self.assertEqual(s, (u'\xff', 'utf-7'))

        s = util.guess_decode_from_terminal(u'\xff'.encode('utf-8'), Term)
        self.assertEqual(s, (u'\xff', 'utf-8'))

    def test_add_metaclass(self):
        class Meta(type):
            pass

        @util.add_metaclass(Meta)
        class Cls:
            pass

        self.assertEqual(type(Cls), Meta)

class ConsoleTest(unittest.TestCase):

    def test_ansiformat(self):
        f = console.ansiformat
        c = console.codes
        all_attrs = f('+*_blue_*+', 'text')
        self.assertTrue(c['blue'] in all_attrs and c['blink'] in all_attrs
                        and c['bold'] in all_attrs and c['underline'] in all_attrs
                        and c['reset'] in all_attrs)
        self.assertRaises(KeyError, f, '*mauve*', 'text')

    def test_functions(self):
        self.assertEqual(console.reset_color(), console.codes['reset'])
        self.assertEqual(console.colorize('blue', 'text'),
                         console.codes['blue'] + 'text' + console.codes['reset'])
<|endoftext|>


<commit_before>import os
from contextlib import contextmanager
from doctest import DocTestSuite
from tempfile import NamedTemporaryFile, TemporaryDirectory
from unittest import TestCase

from runcommands.commands import copy_file

def load_tests(loader, tests, ignore):
    tests.addTests(DocTestSuite("runcommands.util.decorators"))
    tests.addTests(DocTestSuite("runcommands.util.misc"))
    tests.addTests(DocTestSuite("runcommands.util.path"))
    tests.addTests(DocTestSuite("runcommands.util.string"))
    return tests

@contextmanager
def copy(contents, destination_dir=False, **kwargs):
    with NamedTemporaryFile("w", delete=False) as tp:
        tp.write(contents)

    source = tp.name

    if destination_dir:
        with TemporaryDirectory() as destination:
            path = copy_file(source, destination, **kwargs)
            yield source, destination, path
            os.remove(source)
    else:
        destination = source + ".copy"
        path = copy_file(source, destination, **kwargs)
        yield source, destination, path
        os.remove(source)
        os.remove(path)

class TestCopyFileCommand(TestCase):
    def assertDestination(self, source, destination, path, expected_contents):
        self.assertExists(path)
        self.assertNotEqual(source, destination)
        self.assertNotEqual(source, path)
        self.assertContentsEqual(path, expected_contents)
        if not os.path.isdir(destination):
            self.assertTrue(path.endswith(".copy"))

    def assertExists(self, path):
        self.assertTrue(os.path.isfile(path))

    def assertContentsEqual(self, path, expected_contents):
        with open(path) as fp:
            read_contents = fp.read()
        self.assertEqual(read_contents, expected_contents)

    def test_copy(self):
        with copy("xyz") as (source, destination, path):
            self.assertDestination(source, destination, path, "xyz")

    def test_copy_template_format(self):
        with copy("{xyz} ${{xyz}}", template=True, context={"xyz": 123}) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "123 ${xyz}")

    def test_copy_template_string(self):
        with copy("${xyz} $${xyz}", template="string", context={"xyz": 123}) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "123 ${xyz}")

    def test_copy_to_directory(self):
        with copy("xyz", destination_dir=True) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "xyz")
            self.assertEqual(os.path.dirname(path), destination)
            self.assertNotEqual(os.path.dirname(source), os.path.dirname(path))
            self.assertEqual(os.path.basename(source), os.path.basename(path))
<commit_msg>Change how doctests are set up<commit_after>import os
from contextlib import contextmanager
from doctest import DocTestSuite
from tempfile import NamedTemporaryFile, TemporaryDirectory
from unittest import TestCase

import runcommands.util.decorators
import runcommands.util.misc
import runcommands.util.path
import runcommands.util.string
from runcommands.commands import copy_file

def load_tests(loader, tests, ignore):
    tests.addTests(DocTestSuite(runcommands.util.decorators))
    tests.addTests(DocTestSuite(runcommands.util.misc))
    tests.addTests(DocTestSuite(runcommands.util.path))
    tests.addTests(DocTestSuite(runcommands.util.string))
    return tests

@contextmanager
def copy(contents, destination_dir=False, **kwargs):
    with NamedTemporaryFile("w", delete=False) as tp:
        tp.write(contents)

    source = tp.name

    if destination_dir:
        with TemporaryDirectory() as destination:
            path = copy_file(source, destination, **kwargs)
            yield source, destination, path
            os.remove(source)
    else:
        destination = source + ".copy"
        path = copy_file(source, destination, **kwargs)
        yield source, destination, path
        os.remove(source)
        os.remove(path)

class TestCopyFileCommand(TestCase):
    def assertDestination(self, source, destination, path, expected_contents):
        self.assertExists(path)
        self.assertNotEqual(source, destination)
        self.assertNotEqual(source, path)
        self.assertContentsEqual(path, expected_contents)
        if not os.path.isdir(destination):
            self.assertTrue(path.endswith(".copy"))

    def assertExists(self, path):
        self.assertTrue(os.path.isfile(path))

    def assertContentsEqual(self, path, expected_contents):
        with open(path) as fp:
            read_contents = fp.read()
        self.assertEqual(read_contents, expected_contents)

    def test_copy(self):
        with copy("xyz") as (source, destination, path):
            self.assertDestination(source, destination, path, "xyz")

    def test_copy_template_format(self):
        with copy("{xyz} ${{xyz}}", template=True, context={"xyz": 123}) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "123 ${xyz}")

    def test_copy_template_string(self):
        with copy("${xyz} $${xyz}", template="string", context={"xyz": 123}) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "123 ${xyz}")

    def test_copy_to_directory(self):
        with copy("xyz", destination_dir=True) as paths:
            source, destination, path = paths
            self.assertDestination(source, destination, path, "xyz")
            self.assertEqual(os.path.dirname(path), destination)
            self.assertNotEqual(os.path.dirname(source), os.path.dirname(path))
            self.assertEqual(os.path.basename(source), os.path.basename(path))
<|endoftext|>


<commit_before>from nose.tools import set_trace
from argparse import ArgumentParser
import logging
import sys
import time

from config import Configuration
from model import (
    _now,
    InvalidPost,
    Post,
    TIME_FORMAT,
)

class Script(object):
    pass

class BotScript(Script):
    """A script that operates on one or more bots."""

    log = logging.getLogger("Bot script")

    @classmethod
    def parser(cls):
        parser = ArgumentParser()
        parser.add_argument(
            '--config',
            help="Directory containing the botfriend database.",
            required=True,
        )
        parser.add_argument(
            '--bot',
            help='Operate on this bot.',
            nargs='*'
        )
        return parser

    def __init__(self):
        parser = self.parser()
        self.args = parser.parse_args()
        self.config = Configuration.from_directory(self.args.config, self.args.bot)

    def run(self):
        found = False
        for model in self.config.bots:
            try:
                found = True
                self.process_bot(model)
            except InvalidPost, e:
                # This _should_ crash the whole script -- we don't
                # want to commit invalid posts to the database.
                raise e
            except Exception, e:
                # Don't let a 'normal' error crash the whole script.
                model.implementation.log.error(e.message, exc_info=e)
        self.config._db.commit()
        if not found:
            self.log.error("Could not find any bot named %s in %s.",
                           self.args.bot, self.config.directory
            )

    def process_bot(self, bot_model):
        raise NotImplementedError()

class SingleBotScript(BotScript):
    """A script that _must_ be run against a single bot."""

    @classmethod
    def parser(cls):
        parser = ArgumentParser()
        parser.add_argument(
            '--config',
            help="Directory containing the botfriend database.",
            required=True,
        )
        parser.add_argument(
            '--bot',
            help='Operate on this bot.',
            required=True
        )
        return parser

class RepublicationScript(BotScript):
    """Attempt to publish already created posts that failed in their
    delivery.
    """
    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of posts to republish per bot.",
            type=int,
            default=1
        )
        return parser

    def process_bot(self, bot_model):
        undelivered = bot_model.undeliverable_posts.limit(self.args.limit)
        for post in undelivered:
            for publication in post.publications:
                if not publication.error:
                    continue
                # Find the publisher responsible for this
                matches  = [x for x in bot_model.implementation.publishers
                            if x.service == publication.service]
                if not matches:
                    # This bot doesn't use this publisher anymore.
                    continue
                [publisher] = matches
                bot_model.log.info(
                    "Attempting to republish to %s: %s" % (
                        publication.service,
                        post.content
                    )
                )
                publisher.publish(post, publication)
                if publication.error:
                    bot_model.log.info("Failure: %s" % publication.error)
                else:
                    bot_model.log.info("Success!")

class DashboardScript(BotScript):
    """Display the current status of one or more bots."""
    def process_bot(self, bot_model):

        now = _now()
        recent = bot_model.recent_posts().limit(1).all()
        if not recent:
            bot_model.log.info("Has never posted.")
        else:
            [recent] = recent
            bot_model.log.info("Most recent post: %s" % recent.content)
            for publication in recent.publications:
                if publication.error:
                    bot_model.log.info(
                        "%s ERROR: %s" % (publication.service, publication.error)
                    )
                else:
                    bot_model.log.info(
                        "%s posted %dm ago (%s)" % (
                            publication.service,
                            (now-publication.most_recent_attempt).total_seconds()/60,
                            publication.most_recent_attempt,
                        )
                    )

        def announce_list(count, content, what):
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d %s %s" % (count, item, what))
            bot_model.log.info("Next up: %s" % content)

        # Announce scheduled posts.
        scheduled = bot_model.scheduled
        next_post_time = bot_model.next_post_time
        if len(scheduled) > 0:
            first = scheduled[0]
            announce_list(len(scheduled), first.content, "scheduled")
            next_post_time = first.publish_at or bot_model.next_post_time

        # Announce backlog posts.
        try:
            backlog = bot_model.backlog
            if backlog:
                first = backlog[0]
                announce_list(backlog, count, first, "in backlog")
        except ValueError, e:
            pass

        if next_post_time:
            minutes = (next_post_time-now).total_seconds()/60
            if minutes < 0:
                when = "ASAP"
            else:
                when = "in %dm" % minutes
            bot_model.log.info("Next post %s" % when)
        else:
            bot_model.log.info("Next post not scheduled.")

class PostScript(BotScript):
    """Create a new post for one or all bots."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            '--dry-run',
            help="Show what would be posted, but don't post it or commit to the database.",
            action='store_true'
        )
        parser.add_argument(
            '--force',
            help="Post even if the scheduler would not normally post now.",
            action='store_true'
        )
        return parser

    def process_bot(self, bot_model):
        implementation = bot_model.implementation
        if self.args.force:
            bot_model.next_post_time = _now()
        posts = implementation.publishable_posts
        if self.args.dry_run:
            print bot_model.name
            for post in posts:
                print post.content
                print "-" * 80
                return

        # We're doing this for real.
        for post in posts:
            for publication in implementation.publish(post):
                publication.post.bot.log.info(publication.display())
        self.config._db.commit()

class StateShowScript(BotScript):
    """Show the internal state for a bot."""

    def process_bot(self, bot_model):
        last_update = bot_model.last_state_update_time
        print "State for %s (last update %s)" % (bot_model.name, last_update)
        print bot_model.state

class StateSetScript(SingleBotScript):
    """Set the internal state for a bot."""

    @classmethod
    def parser(cls):
        parser = SingleBotScript.parser()
        parser.add_argument(
            "--file",
            help="Load from this file instead of standard input.",
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        if self.args.file:
            fh = open(self.args.file)
        else:
            fh = sys.stdin
        data = fh.read().decode("utf8")
        bot_model.implementation.set_state(data)
        print bot_model.state

class StateRefreshScript(BotScript):
    """Refresh the internal state for a bot."""

    def process_bot(self, bot_model):
        bot_model.implementation.check_and_update_state(force=True)
        print "State for %s (last update %s)" % (
            bot_model.name, bot_model.last_state_update_time
        )
        print bot_model.state

class StressTestScript(BotScript):
    """Stress-test a bot's generative capabilities without posting anything."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            '--rounds',
            help="Run the bot's generator this many times. (Default is 10,000)",
            type=int,
            default=10000
        )
        return parser

    def process_bot(self, bot_model):
        bot_model.implementation.stress_test(self.args.rounds)

class PublisherTestScript(BotScript):
    """Verify  that a bot's publishers are functioning without posting anything."""

    def process_bot(self, bot_model):
        for publisher in bot_model.implementation.publishers:
            try:
                result = publisher.self_test() or ""
                print "GOOD %s %s %s" % (bot_model.name, publisher.service, result)
            except Exception, e:
                print "FAIL %s %s: %s" % (
                    bot_model.name, publisher.service, e
                )

class BacklogShowScript(BotScript):
    """Show the backlog posts for a bot."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of backlog posts shown.",
            type=int,
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        backlog = bot_model.implementation.backlog
        count = len(backlog)
        if self.args.limit:
            max_i = self.args.limit - 1
        else:
            max_i = None
        if count:
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d %s in backlog" % (count, item))
            for i, content in enumerate(backlog):
                bot_model.log.info(content)
                if i > max_i:
                    break
        else:
            bot_model.log.info("No backlog.")

class BacklogLoadScript(SingleBotScript):

    @classmethod
    def parser(cls):
        parser = SingleBotScript.parser()
        parser.add_argument(
            "--file",
            help="Load from this file instead of standard input.",
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        if self.args.file:
            fh = open(self.args.file)
        else:
            fh = sys.stdin
        # Process one backlog item per line of the input file.
        items = [x.strip().decode("utf8") for x in fh.readlines()]
        bot_model.implementation.extend_backlog(items)
        self.log.info("Appended %d items to backlog." % len(items))
        self.log.info("Backlog size now %d items" % len(bot_model.backlog))

class BacklogClearScript(SingleBotScript):

    def process_bot(self, bot_model):
        backlog = bot_model.backlog
        if backlog:
            bot_model.log.warn(
                "About to clear the backlog for %s.", bot_model.name
            )
            bot_model.log.warn(
                "Sleeping for 2 seconds to give you a chance to Ctrl-C."
            )
            time.sleep(2)
            bot_model.implementation.clear_backlog()

class ScheduledPostsShowScript(BotScript):
    """Show the scheduled posts for a bot."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of scheduled posts shown.",
            type=int,
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        scheduled = bot_model.scheduled
        count = scheduled.count()
        if self.args.limit:
            max_i = self.args.limit - 1
        else:
            max_i = None
        if count:
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d scheduled %s" % (count, item))
            for i, post in enumerate(scheduled):
                if max_i is not None and i > max_i:
                    break
                if post.publish_at:
                    when_post = post.publish_at.strftime(TIME_FORMAT)
                elif i == 0 and bot_model.next_post_time:
                    when_post = bot_model.next_post_time.strftime(TIME_FORMAT)
                else:
                    when_post = "No scheduled time"
                bot_model.log.info("%s | %s" % (when_post, post.content))
        else:
            bot_model.log.info("No scheduled posts.")

class SchedulePostsLoadScript(SingleBotScript):

    def process_bot(self, bot_model):
        bot.schedule_posts()

class ScheduledPostsClearScript(SingleBotScript):

    def process_bot(self, bot_model):
        scheduled = list(bot_model.scheduled)
        if scheduled:
            bot_model.log.warn(
                "About to remove all scheduled posts for %s.", bot_model.name
            )
            bot_model.log.warn(
                "Sleeping for 2 seconds to give you a chance to Ctrl-C."
            )
            time.sleep(2)
            for post in scheduled:
                self.config._db.delete(post)

        # Also reset the next post time.
        bot_model.next_post_time = bot_model.implementation.schedule_next_post()
        if bot_model.next_post_time:
            bot_model.log.info("Next post at %s", bot_model.next_post_time)
        else:
            bot_model.log.info("Ready to post.")
<commit_msg>scheduled is a list now, not a query.<commit_after>from nose.tools import set_trace
from argparse import ArgumentParser
import logging
import sys
import time

from config import Configuration
from model import (
    _now,
    InvalidPost,
    Post,
    TIME_FORMAT,
)

class Script(object):
    pass

class BotScript(Script):
    """A script that operates on one or more bots."""

    log = logging.getLogger("Bot script")

    @classmethod
    def parser(cls):
        parser = ArgumentParser()
        parser.add_argument(
            '--config',
            help="Directory containing the botfriend database.",
            required=True,
        )
        parser.add_argument(
            '--bot',
            help='Operate on this bot.',
            nargs='*'
        )
        return parser

    def __init__(self):
        parser = self.parser()
        self.args = parser.parse_args()
        self.config = Configuration.from_directory(self.args.config, self.args.bot)

    def run(self):
        found = False
        for model in self.config.bots:
            try:
                found = True
                self.process_bot(model)
            except InvalidPost, e:
                # This _should_ crash the whole script -- we don't
                # want to commit invalid posts to the database.
                raise e
            except Exception, e:
                # Don't let a 'normal' error crash the whole script.
                model.implementation.log.error(e.message, exc_info=e)
        self.config._db.commit()
        if not found:
            self.log.error("Could not find any bot named %s in %s.",
                           self.args.bot, self.config.directory
            )

    def process_bot(self, bot_model):
        raise NotImplementedError()

class SingleBotScript(BotScript):
    """A script that _must_ be run against a single bot."""

    @classmethod
    def parser(cls):
        parser = ArgumentParser()
        parser.add_argument(
            '--config',
            help="Directory containing the botfriend database.",
            required=True,
        )
        parser.add_argument(
            '--bot',
            help='Operate on this bot.',
            required=True
        )
        return parser

class RepublicationScript(BotScript):
    """Attempt to publish already created posts that failed in their
    delivery.
    """
    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of posts to republish per bot.",
            type=int,
            default=1
        )
        return parser

    def process_bot(self, bot_model):
        undelivered = bot_model.undeliverable_posts.limit(self.args.limit)
        for post in undelivered:
            for publication in post.publications:
                if not publication.error:
                    continue
                # Find the publisher responsible for this
                matches  = [x for x in bot_model.implementation.publishers
                            if x.service == publication.service]
                if not matches:
                    # This bot doesn't use this publisher anymore.
                    continue
                [publisher] = matches
                bot_model.log.info(
                    "Attempting to republish to %s: %s" % (
                        publication.service,
                        post.content
                    )
                )
                publisher.publish(post, publication)
                if publication.error:
                    bot_model.log.info("Failure: %s" % publication.error)
                else:
                    bot_model.log.info("Success!")

class DashboardScript(BotScript):
    """Display the current status of one or more bots."""
    def process_bot(self, bot_model):

        now = _now()
        recent = bot_model.recent_posts().limit(1).all()
        if not recent:
            bot_model.log.info("Has never posted.")
        else:
            [recent] = recent
            bot_model.log.info("Most recent post: %s" % recent.content)
            for publication in recent.publications:
                if publication.error:
                    bot_model.log.info(
                        "%s ERROR: %s" % (publication.service, publication.error)
                    )
                else:
                    bot_model.log.info(
                        "%s posted %dm ago (%s)" % (
                            publication.service,
                            (now-publication.most_recent_attempt).total_seconds()/60,
                            publication.most_recent_attempt,
                        )
                    )

        def announce_list(count, content, what):
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d %s %s" % (count, item, what))
            bot_model.log.info("Next up: %s" % content)

        # Announce scheduled posts.
        scheduled = bot_model.scheduled
        next_post_time = bot_model.next_post_time
        if len(scheduled) > 0:
            first = scheduled[0]
            announce_list(len(scheduled), first.content, "scheduled")
            next_post_time = first.publish_at or bot_model.next_post_time

        # Announce backlog posts.
        try:
            backlog = bot_model.backlog
            if backlog:
                first = backlog[0]
                announce_list(backlog, count, first, "in backlog")
        except ValueError, e:
            pass

        if next_post_time:
            minutes = (next_post_time-now).total_seconds()/60
            if minutes < 0:
                when = "ASAP"
            else:
                when = "in %dm" % minutes
            bot_model.log.info("Next post %s" % when)
        else:
            bot_model.log.info("Next post not scheduled.")

class PostScript(BotScript):
    """Create a new post for one or all bots."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            '--dry-run',
            help="Show what would be posted, but don't post it or commit to the database.",
            action='store_true'
        )
        parser.add_argument(
            '--force',
            help="Post even if the scheduler would not normally post now.",
            action='store_true'
        )
        return parser

    def process_bot(self, bot_model):
        implementation = bot_model.implementation
        if self.args.force:
            bot_model.next_post_time = _now()
        posts = implementation.publishable_posts
        if self.args.dry_run:
            print bot_model.name
            for post in posts:
                print post.content
                print "-" * 80
                return

        # We're doing this for real.
        for post in posts:
            for publication in implementation.publish(post):
                publication.post.bot.log.info(publication.display())
        self.config._db.commit()

class StateShowScript(BotScript):
    """Show the internal state for a bot."""

    def process_bot(self, bot_model):
        last_update = bot_model.last_state_update_time
        print "State for %s (last update %s)" % (bot_model.name, last_update)
        print bot_model.state

class StateSetScript(SingleBotScript):
    """Set the internal state for a bot."""

    @classmethod
    def parser(cls):
        parser = SingleBotScript.parser()
        parser.add_argument(
            "--file",
            help="Load from this file instead of standard input.",
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        if self.args.file:
            fh = open(self.args.file)
        else:
            fh = sys.stdin
        data = fh.read().decode("utf8")
        bot_model.implementation.set_state(data)
        print bot_model.state

class StateRefreshScript(BotScript):
    """Refresh the internal state for a bot."""

    def process_bot(self, bot_model):
        bot_model.implementation.check_and_update_state(force=True)
        print "State for %s (last update %s)" % (
            bot_model.name, bot_model.last_state_update_time
        )
        print bot_model.state

class StressTestScript(BotScript):
    """Stress-test a bot's generative capabilities without posting anything."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            '--rounds',
            help="Run the bot's generator this many times. (Default is 10,000)",
            type=int,
            default=10000
        )
        return parser

    def process_bot(self, bot_model):
        bot_model.implementation.stress_test(self.args.rounds)

class PublisherTestScript(BotScript):
    """Verify  that a bot's publishers are functioning without posting anything."""

    def process_bot(self, bot_model):
        for publisher in bot_model.implementation.publishers:
            try:
                result = publisher.self_test() or ""
                print "GOOD %s %s %s" % (bot_model.name, publisher.service, result)
            except Exception, e:
                print "FAIL %s %s: %s" % (
                    bot_model.name, publisher.service, e
                )

class BacklogShowScript(BotScript):
    """Show the backlog posts for a bot."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of backlog posts shown.",
            type=int,
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        backlog = bot_model.implementation.backlog
        count = len(backlog)
        if self.args.limit:
            max_i = self.args.limit - 1
        else:
            max_i = None
        if count:
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d %s in backlog" % (count, item))
            for i, content in enumerate(backlog):
                bot_model.log.info(content)
                if i > max_i:
                    break
        else:
            bot_model.log.info("No backlog.")

class BacklogLoadScript(SingleBotScript):

    @classmethod
    def parser(cls):
        parser = SingleBotScript.parser()
        parser.add_argument(
            "--file",
            help="Load from this file instead of standard input.",
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        if self.args.file:
            fh = open(self.args.file)
        else:
            fh = sys.stdin
        # Process one backlog item per line of the input file.
        items = [x.strip().decode("utf8") for x in fh.readlines()]
        bot_model.implementation.extend_backlog(items)
        self.log.info("Appended %d items to backlog." % len(items))
        self.log.info("Backlog size now %d items" % len(bot_model.backlog))

class BacklogClearScript(SingleBotScript):

    def process_bot(self, bot_model):
        backlog = bot_model.backlog
        if backlog:
            bot_model.log.warn(
                "About to clear the backlog for %s.", bot_model.name
            )
            bot_model.log.warn(
                "Sleeping for 2 seconds to give you a chance to Ctrl-C."
            )
            time.sleep(2)
            bot_model.implementation.clear_backlog()

class ScheduledPostsShowScript(BotScript):
    """Show the scheduled posts for a bot."""

    @classmethod
    def parser(cls):
        parser = BotScript.parser()
        parser.add_argument(
            "--limit",
            help="Limit the number of scheduled posts shown.",
            type=int,
            default=None
        )
        return parser

    def process_bot(self, bot_model):
        scheduled = bot_model.scheduled
        count = len(scheduled)
        if self.args.limit:
            max_i = self.args.limit - 1
        else:
            max_i = None
        if count:
            if count == 1:
                item = "post"
            else:
                item = "posts"
            bot_model.log.info("%d scheduled %s" % (count, item))
            for i, post in enumerate(scheduled):
                if max_i is not None and i > max_i:
                    break
                if post.publish_at:
                    when_post = post.publish_at.strftime(TIME_FORMAT)
                elif i == 0 and bot_model.next_post_time:
                    when_post = bot_model.next_post_time.strftime(TIME_FORMAT)
                else:
                    when_post = "No scheduled time"
                bot_model.log.info("%s | %s" % (when_post, post.content))
        else:
            bot_model.log.info("No scheduled posts.")

class SchedulePostsLoadScript(SingleBotScript):

    def process_bot(self, bot_model):
        bot.schedule_posts()

class ScheduledPostsClearScript(SingleBotScript):

    def process_bot(self, bot_model):
        scheduled = list(bot_model.scheduled)
        if scheduled:
            bot_model.log.warn(
                "About to remove all scheduled posts for %s.", bot_model.name
            )
            bot_model.log.warn(
                "Sleeping for 2 seconds to give you a chance to Ctrl-C."
            )
            time.sleep(2)
            for post in scheduled:
                self.config._db.delete(post)

        # Also reset the next post time.
        bot_model.next_post_time = bot_model.implementation.schedule_next_post()
        if bot_model.next_post_time:
            bot_model.log.info("Next post at %s", bot_model.next_post_time)
        else:
            bot_model.log.info("Ready to post.")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import unittest

import datetime
import os
import re
import sys
import pytz

import when

sys.path.insert(0, os.path.abspath('..'))
sys.path.append('.')

class WhenTest(unittest.TestCase):
    def setUp(self):
        when.unset_utc()

        self.one_day = datetime.timedelta(days=1)
        self.one_second = datetime.timedelta(seconds=1)

        self.today = datetime.date.today()

        self.now = datetime.datetime.now()
        self.utc = datetime.datetime.utcnow()

        env_timezone = os.getenv('TIMEZONE')

        if env_timezone:
            self.timezone = env_timezone
        elif os.path.exists('/etc/localtime'):
            localtime_path = os.path.realpath('/etc/localtime')
            self.timezone = re.findall('([^/]*/[^/]*)$', localtime_path)[0]
        else:
            self.timezone = 'America/New_York'

    def test__add_time(self):
        # Test change between months with dfferent number of days
        test_value = datetime.datetime(2012, 3, 31)

        expected_value = datetime.datetime(2012, 5, 1)
        result = when._add_time(test_value, months=1)
        self.assertEqual(result, expected_value)

        # Test values going back into February of a leap year
        expected_value = datetime.datetime(2012, 3, 2)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        test_value = datetime.datetime(2012, 3, 30)

        expected_value = datetime.datetime(2012, 3, 1)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        test_value = datetime.datetime(2011, 3, 31)

        expected_value = datetime.datetime(2011, 3, 3)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        # Test leap day specifically
        test_value = datetime.datetime(2012, 2, 29)

        expected_value = datetime.datetime(2013, 3, 1)
        result = when._add_time(test_value, years=1)
        self.assertEqual(result, expected_value)

        expected_value = datetime.datetime(2011, 3, 1)
        result = when._add_time(test_value, years=-1)
        self.assertEqual(result, expected_value)

    def test__add_time_assert(self):
        self.assertRaises(AssertionError, when._add_time, 'a')

    def test__is_date_type(self):
        self.assertFalse(when._is_date_type('a'))
        self.assertFalse(when._is_date_type(1))
        self.assertFalse(when._is_date_type(['a']))

        self.assertTrue(when._is_date_type(self.today))
        self.assertTrue(when._is_date_type(self.now))
        self.assertTrue(when._is_date_type(self.now.time()))

    def test_ever(self):
        old_result = None
        for i in range(50):
            result = when.ever()
            self.assertTrue(isinstance(result, datetime.datetime))
            self.assertNotEqual(result, old_result)
            old_result = result

    def test_format(self):
        now = when.now()
        today = when.today()
        current_time = now.time()

        for format_string in ('%a', '%A', '%b', '%B', '%c', '%d', '%f', '%H',
                              '%I', '%j', '%m', '%M', '%p', '%S', '%U', '%w',
                              '%W', '%x', '%X', '%y', '%Y', '%z', '%Z',
                              '%A, %B %d, %Y %I:%M %p'):
            # Test date objects
            builtin_date = now.strftime(format_string)
            result_date = when.format(now, format_string)
            self.assertEqual(builtin_date, result_date)

            # Test datetime objects
            builtin_datetime = today.strftime(format_string)
            result_datetime = when.format(today, format_string)
            self.assertEqual(builtin_datetime, result_datetime)

            # Test time objects
            builtin_time = current_time.strftime(format_string)
            result_time = when.format(current_time, format_string)
            self.assertEqual(builtin_time, result_time)

    def test_format_assert(self):
        self.assertRaises(AssertionError, when.format, 'a', '%a')

    def test_is_timezone_aware(self):
        naive = when.now()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_aware(aware))
        self.assertFalse(when.is_timezone_aware(naive))

        naive = naive.time()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_aware(aware))
        self.assertFalse(when.is_timezone_aware(naive))

    def test_is_timezone_aware_assert(self):
        today = when.today()
        self.assertRaises(AssertionError, when.is_timezone_aware, today)

    def test_is_timezone_naive(self):
        naive = when.now()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_naive(naive))
        self.assertFalse(when.is_timezone_naive(aware))

        naive = naive.time()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_naive(naive))
        self.assertFalse(when.is_timezone_naive(aware))

    def test_is_timezone_aware_naive(self):
        today = when.today()
        self.assertRaises(AssertionError, when.is_timezone_aware, today)

    def test_now(self):
        now = when.now()
        utc = when.now(True)

        # Unfortunately the clock keeps ticking each time we capture a value
        # for now so we can't do a direct comparison with assertEqual.
        # It's probably safe to assume the now function is working as long as
        # the difference is less than a second. There's probably a better way
        # to test this, but for now it's sufficient.
        self.assertTrue(now - self.now < self.one_second)
        self.assertTrue(utc - self.utc < self.one_second)

    def test_set_utc(self):
        when.set_utc()
        self.assertEqual(when._FORCE_UTC, True)

    def test_shift(self):
        first = when.shift(self.utc, from_tz='UTC', to_tz='America/New_York')
        second = when.shift(first, from_tz='America/New_York', to_tz='UTC')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

        # Local time
        if self.timezone == 'UTC':
            # This block is needed for tests run in an environment set to UTC.
            first = when.shift(self.now, to_tz='America/New_York')
            second = when.shift(first, from_tz='America/New_York')
        else:
            first = when.shift(self.now, to_tz='UTC')
            second = when.shift(first, from_tz='UTC')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.now)
        self.assertEqual(second, self.now)

        # Set utc parameter to true
        first = when.shift(self.utc, to_tz='America/New_York', utc=True)
        second = when.shift(first, from_tz='America/New_York', utc=True)

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

        # Force UTC
        when.set_utc()
        first = when.shift(self.utc, to_tz='America/New_York')
        second = when.shift(first, from_tz='America/New_York')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

    def test_shift_assert(self):
        self.assertRaises(AssertionError, when.shift, 'a')
        self.assertRaises(AssertionError, when.shift, when.today())

    def test_shift_aware(self):
        central = pytz.timezone('America/Chicago')

        now_aware = central.localize(self.now)

        # Make sure the datetime's time zone is the respected
        first = when.shift(now_aware, to_tz='America/New_York')
        second = when.shift(self.now, from_tz='America/Chicago', to_tz='America/New_York')

        self.assertEqual(first, second)

        # Also make sure the from_tz parameter is ignored
        first = when.shift(now_aware, from_tz='UTC', to_tz='America/New_York')

        self.assertEqual(first, second)

        # Also make sure the utc parameter is ignored
        first = when.shift(now_aware, to_tz='America/New_York', utc=True)

        self.assertEqual(first, second)

    def test_timezone(self):
        self.assertEqual(when.timezone(), self.timezone)

    def test_timezone_object(self):
        local_timezone = pytz.timezone(self.timezone)
        self.assertEqual(when.timezone_object(), local_timezone)

    def test_timezones(self):
        # Make sure all_timezones() matches pytz's versions
        all_timezones = when.all_timezones()
        self.assertEqual(all_timezones, pytz.all_timezones)
        all_timezones_set = when.all_timezones_set()
        self.assertEqual(all_timezones_set, pytz.all_timezones_set)

        # Make sure common_timezones() matches pytz's versions
        common_timezones = when.common_timezones()
        self.assertEqual(common_timezones, pytz.common_timezones)
        common_timezones_set = when.common_timezones_set()
        self.assertEqual(common_timezones_set, pytz.common_timezones_set)

    def test_today(self):
        self.assertEqual(when.today(), self.today)

    def test_tomorrow(self):
        self.assertEqual(when.tomorrow(), self.today + self.one_day)

    def test_unset_utc(self):
        when.unset_utc()
        self.assertEqual(when._FORCE_UTC, False)

    def test_yesterday(self):
        self.assertEqual(when.yesterday(), self.today - self.one_day)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Split test_timezones() into multiple tests<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import unittest

import datetime
import os
import re
import sys
import pytz

import when

sys.path.insert(0, os.path.abspath('..'))
sys.path.append('.')

class WhenTest(unittest.TestCase):
    def setUp(self):
        when.unset_utc()

        self.one_day = datetime.timedelta(days=1)
        self.one_second = datetime.timedelta(seconds=1)

        self.today = datetime.date.today()

        self.now = datetime.datetime.now()
        self.utc = datetime.datetime.utcnow()

        env_timezone = os.getenv('TIMEZONE')

        if env_timezone:
            self.timezone = env_timezone
        elif os.path.exists('/etc/localtime'):
            localtime_path = os.path.realpath('/etc/localtime')
            self.timezone = re.findall('([^/]*/[^/]*)$', localtime_path)[0]
        else:
            self.timezone = 'America/New_York'

    def test__add_time(self):
        # Test change between months with dfferent number of days
        test_value = datetime.datetime(2012, 3, 31)

        expected_value = datetime.datetime(2012, 5, 1)
        result = when._add_time(test_value, months=1)
        self.assertEqual(result, expected_value)

        # Test values going back into February of a leap year
        expected_value = datetime.datetime(2012, 3, 2)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        test_value = datetime.datetime(2012, 3, 30)

        expected_value = datetime.datetime(2012, 3, 1)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        test_value = datetime.datetime(2011, 3, 31)

        expected_value = datetime.datetime(2011, 3, 3)
        result = when._add_time(test_value, months=-1)
        self.assertEqual(result, expected_value)

        # Test leap day specifically
        test_value = datetime.datetime(2012, 2, 29)

        expected_value = datetime.datetime(2013, 3, 1)
        result = when._add_time(test_value, years=1)
        self.assertEqual(result, expected_value)

        expected_value = datetime.datetime(2011, 3, 1)
        result = when._add_time(test_value, years=-1)
        self.assertEqual(result, expected_value)

    def test__add_time_assert(self):
        self.assertRaises(AssertionError, when._add_time, 'a')

    def test__is_date_type(self):
        self.assertFalse(when._is_date_type('a'))
        self.assertFalse(when._is_date_type(1))
        self.assertFalse(when._is_date_type(['a']))

        self.assertTrue(when._is_date_type(self.today))
        self.assertTrue(when._is_date_type(self.now))
        self.assertTrue(when._is_date_type(self.now.time()))

    def test_all_timezones(self):
        # Make sure all_timezones() matches pytz's version
        all_timezones = when.all_timezones()
        self.assertEqual(all_timezones, pytz.all_timezones)

    def test_all_timezones_set(self):
        # Make sure all_timezones_set() matches pytz's version
        all_timezones_set = when.all_timezones_set()
        self.assertEqual(all_timezones_set, pytz.all_timezones_set)

    def test_common_timezones(self):
        # Make sure common_timezones() matches pytz's version
        common_timezones = when.common_timezones()
        self.assertEqual(common_timezones, pytz.common_timezones)

    def test_common_timezones_set(self):
        # Make sure common_timezones_set() matches pytz's version
        common_timezones_set = when.common_timezones_set()
        self.assertEqual(common_timezones_set, pytz.common_timezones_set)

    def test_ever(self):
        old_result = None
        for i in range(50):
            result = when.ever()
            self.assertTrue(isinstance(result, datetime.datetime))
            self.assertNotEqual(result, old_result)
            old_result = result

    def test_format(self):
        now = when.now()
        today = when.today()
        current_time = now.time()

        for format_string in ('%a', '%A', '%b', '%B', '%c', '%d', '%f', '%H',
                              '%I', '%j', '%m', '%M', '%p', '%S', '%U', '%w',
                              '%W', '%x', '%X', '%y', '%Y', '%z', '%Z',
                              '%A, %B %d, %Y %I:%M %p'):
            # Test date objects
            builtin_date = now.strftime(format_string)
            result_date = when.format(now, format_string)
            self.assertEqual(builtin_date, result_date)

            # Test datetime objects
            builtin_datetime = today.strftime(format_string)
            result_datetime = when.format(today, format_string)
            self.assertEqual(builtin_datetime, result_datetime)

            # Test time objects
            builtin_time = current_time.strftime(format_string)
            result_time = when.format(current_time, format_string)
            self.assertEqual(builtin_time, result_time)

    def test_format_assert(self):
        self.assertRaises(AssertionError, when.format, 'a', '%a')

    def test_is_timezone_aware(self):
        naive = when.now()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_aware(aware))
        self.assertFalse(when.is_timezone_aware(naive))

        naive = naive.time()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_aware(aware))
        self.assertFalse(when.is_timezone_aware(naive))

    def test_is_timezone_aware_assert(self):
        today = when.today()
        self.assertRaises(AssertionError, when.is_timezone_aware, today)

    def test_is_timezone_naive(self):
        naive = when.now()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_naive(naive))
        self.assertFalse(when.is_timezone_naive(aware))

        naive = naive.time()
        aware = naive.replace(tzinfo=pytz.UTC)

        self.assertTrue(when.is_timezone_naive(naive))
        self.assertFalse(when.is_timezone_naive(aware))

    def test_is_timezone_aware_naive(self):
        today = when.today()
        self.assertRaises(AssertionError, when.is_timezone_aware, today)

    def test_now(self):
        now = when.now()
        utc = when.now(True)

        # Unfortunately the clock keeps ticking each time we capture a value
        # for now so we can't do a direct comparison with assertEqual.
        # It's probably safe to assume the now function is working as long as
        # the difference is less than a second. There's probably a better way
        # to test this, but for now it's sufficient.
        self.assertTrue(now - self.now < self.one_second)
        self.assertTrue(utc - self.utc < self.one_second)

    def test_set_utc(self):
        when.set_utc()
        self.assertEqual(when._FORCE_UTC, True)

    def test_shift(self):
        first = when.shift(self.utc, from_tz='UTC', to_tz='America/New_York')
        second = when.shift(first, from_tz='America/New_York', to_tz='UTC')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

        # Local time
        if self.timezone == 'UTC':
            # This block is needed for tests run in an environment set to UTC.
            first = when.shift(self.now, to_tz='America/New_York')
            second = when.shift(first, from_tz='America/New_York')
        else:
            first = when.shift(self.now, to_tz='UTC')
            second = when.shift(first, from_tz='UTC')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.now)
        self.assertEqual(second, self.now)

        # Set utc parameter to true
        first = when.shift(self.utc, to_tz='America/New_York', utc=True)
        second = when.shift(first, from_tz='America/New_York', utc=True)

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

        # Force UTC
        when.set_utc()
        first = when.shift(self.utc, to_tz='America/New_York')
        second = when.shift(first, from_tz='America/New_York')

        self.assertNotEqual(first, second)
        self.assertNotEqual(first, self.utc)
        self.assertEqual(second, self.utc)

    def test_shift_assert(self):
        self.assertRaises(AssertionError, when.shift, 'a')
        self.assertRaises(AssertionError, when.shift, when.today())

    def test_shift_aware(self):
        central = pytz.timezone('America/Chicago')

        now_aware = central.localize(self.now)

        # Make sure the datetime's time zone is the respected
        first = when.shift(now_aware, to_tz='America/New_York')
        second = when.shift(self.now, from_tz='America/Chicago', to_tz='America/New_York')

        self.assertEqual(first, second)

        # Also make sure the from_tz parameter is ignored
        first = when.shift(now_aware, from_tz='UTC', to_tz='America/New_York')

        self.assertEqual(first, second)

        # Also make sure the utc parameter is ignored
        first = when.shift(now_aware, to_tz='America/New_York', utc=True)

        self.assertEqual(first, second)

    def test_timezone(self):
        self.assertEqual(when.timezone(), self.timezone)

    def test_timezone_object(self):
        local_timezone = pytz.timezone(self.timezone)
        self.assertEqual(when.timezone_object(), local_timezone)

    def test_today(self):
        self.assertEqual(when.today(), self.today)

    def test_tomorrow(self):
        self.assertEqual(when.tomorrow(), self.today + self.one_day)

    def test_unset_utc(self):
        when.unset_utc()
        self.assertEqual(when._FORCE_UTC, False)

    def test_yesterday(self):
        self.assertEqual(when.yesterday(), self.today - self.one_day)

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


<commit_before># -*- coding: utf-8 -*-
import os.path
import unittest
try:
    from cStringIO import StringIO as BytesIO
except ImportError:
    from io import BytesIO

import fiole
from ._common import (PY3, ENVIRON, handle_single_request,
                      FORM_DATA_CONTENT_TYPE, FORM_DATA_1, FORM_DATA_2)
u = lambda s: s.decode('utf-8') if isinstance(s, bytes) else s
b = lambda s: s if isinstance(s, bytes) else s.encode('utf-8')

class FioleTestCase(unittest.TestCase):
    _cookie_secret = fiole.SECRET_KEY

    def setUp(self):
        fiole.Fiole.push()
        fiole.SECRET_KEY = self._cookie_secret

    def tearDown(self):
        fiole.Fiole.pop()
        fiole.SECRET_KEY = self._cookie_secret

    def assertNoError(self, response):
        self.assertFalse(response['errors'], msg=response['errors'])
        self.assertEqual(response['status'], '200 OK')

    def test_simple(self):
        @fiole.get('/')
        def index(request):
            return 'Hello World!'
        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('Hello World!')])
        self.assertEqual(rv, {
            'status': '200 OK',
            'headers': [('Content-Type', 'text/html; charset=utf-8'),
                        ('Content-Length', '12')],
            'data': [b('Hello World!')],
            'errors': '',
        })

    def test_request_dispatching(self):
        @fiole.get('/')
        def index(request):
            return request.method

        @fiole.get('/more')
        @fiole.post('/more')
        def more(request):
            return request.method

        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('GET')])

        rv = handle_single_request('POST /')
        self.assertEqual(rv['status'], '405 Method Not Allowed')

        rv = handle_single_request('HEAD /')
        self.assertNoError(rv)
        self.assertFalse(rv['data'])

        rv = handle_single_request('POST /more')
        self.assertEqual(rv['data'], [b('POST')])

        rv = handle_single_request('GET /more')
        self.assertEqual(rv['data'], [b('GET')])

        rv = handle_single_request('DELETE /more')
        self.assertEqual(rv['status'], '405 Method Not Allowed')

    def test_request_route_placeholder(self):
        @fiole.route(u('/<_mot_>/à/<v_42>'))
        def route1(request, _mot_, v_42):
            return u('%s à %s') % (_mot_, v_42)

        rv = handle_single_request(u('GET /table/à/manger'))
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('table à manger')])

    def test_request_route_regex(self):
        @fiole.route(u('/(?P<a>\d+)/plus/(?P<b>\d+)'))
        def addition(request, a, b):
            (a, b) = int(a), int(b)
            return '%d + %d = %d' % (a, b, a + b)

        rv = handle_single_request(u('GET /25/plus/17/'))
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('25 + 17 = 42')])

        rv = handle_single_request(u('GET /navet/plus/tomate/'))
        self.assertEqual(rv['status'], '404 Not Found')

    def test_error_handling(self):
        @fiole.errorhandler(404)
        def not_found(exc):
            return 'not found'

        @fiole.errorhandler(500)
        def internal_server_error(exc):
            return 'internal server error'

        @fiole.get('/')
        def index(request):
            raise fiole.NotFound('Lost')

        @fiole.get('/error')
        def error(request):
            1 // 0

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '404 Not Found')
        self.assertEqual(rv['data'], [b('not found')])
        self.assertFalse(rv['errors'])
        rv = handle_single_request('GET /error')
        self.assertEqual(rv['status'], '500 Internal Server Error')
        self.assertEqual(rv['data'], [b('internal server error')])
        self.assertIn('ZeroDivisionError', rv['errors'])

    def test_user_error_handling(self):
        class MyException(Exception):
            pass

        @fiole.errorhandler(500)
        def system_error(exc):
            if isinstance(exc, MyException):
                return '42'
            return 'internal server error'

        @fiole.get('/')
        def index(request):
            raise MyException()
        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '500 Internal Server Error')
        self.assertEqual(rv['data'], [b('42')])
        self.assertIn('MyException', rv['errors'])

    def test_response(self):

        @fiole.get('/unicode')
        def from_unicode(request):
            return u('Hällo Wörld')

        @fiole.get('/string')
        def from_string(request):
            return b('Hällo Wörld')

        @fiole.get('/args')
        def from_response(request):
            return fiole.Response('Meh', headers=[('X-Foo', 'Testing')],
                                  status=400, content_type='text/plain')

        rv = handle_single_request('GET /unicode')
        self.assertEqual(rv['data'], [b('Hällo Wörld')])
        rv = handle_single_request('GET /string')
        self.assertEqual(rv['data'], [b('Hällo Wörld')])
        rv = handle_single_request('GET /args')
        self.assertEqual(rv, {
            'status': '400 Bad Request',
            'headers': [('X-Foo', 'Testing'),
                        ('Content-Type', 'text/plain; charset=utf-8'),
                        ('Content-Length', '3')],
            'data': [b('Meh')],
            'errors': '',
        })

        rv = fiole.Response('')
        html_utf8 = 'text/html; charset=utf-8'
        self.assertEqual(rv.status, 200)
        self.assertEqual(rv.output, '')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

        rv = fiole.Response('Awesome')
        self.assertEqual(rv.status, 200)
        self.assertEqual(rv.output, 'Awesome')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

        rv = fiole.Response('W00t', status=404)
        self.assertEqual(rv.status, 404)
        self.assertEqual(rv.output, 'W00t')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

    def test_none_response(self):

        @fiole.get('/')
        def index(request):
            return

        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('')])

    def test_redirect(self):

        @fiole.get('/')
        def index(request):
            raise fiole.Redirect('/far')

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '302 Found')
        self.assertEqual(rv['data'], [b('')])
        self.assertIn(('Location', '/far'), rv['headers'])

    def test_redirect_error_handler(self):

        @fiole.errorhandler(500)
        def handle_500_redirect(request):
            raise fiole.Redirect('/far')

        @fiole.get('/')
        def index(request):
            1 // 0

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '302 Found')
        self.assertEqual(rv['data'], [b('')])
        self.assertIn(('Location', '/far'), rv['headers'])
        # Is it correct?
        self.assertIn('ZeroDivisionError', rv['errors'])

    def test_form_data_1(self):
        environ = dict(ENVIRON)
        environ.update({
            'REQUEST_METHOD': 'POST',
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'CONTENT_LENGTH': len(FORM_DATA_1),
            'wsgi.input': BytesIO(FORM_DATA_1),
        })
        request = fiole.Request(environ)
        self.assertFalse(request.GET)
        self.assertFalse(request.cookies)
        self.assertEqual(sorted(request.POST), ['files', 'submit-name'])
        self.assertEqual(request.POST['submit-name'], 'Larry')
        self.assertEqual(request.POST['files'].filename, 'file1.txt')
        self.assertEqual(request.POST['files'].file.read(),
                         b('... contents of file1.txt ...'))
        self.assertEqual(request.POST['files'].value,
                         b('... contents of file1.txt ...'))
        self.assertEqual(request.POST['files'].type, 'text/plain')
        self.assertEqual(request.POST['files'].disposition, 'form-data')

    def test_form_data_2(self):
        environ = dict(ENVIRON)
        environ.update({
            'REQUEST_METHOD': 'POST',
            'CONTENT_LENGTH': len(FORM_DATA_2),
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'wsgi.input': BytesIO(FORM_DATA_2),
        })
        request = fiole.Request(environ)
        self.assertFalse(request.GET)
        self.assertFalse(request.cookies)
        self.assertEqual(sorted(request.POST), ['files', 'submit-name'])
        self.assertEqual(request.POST['submit-name'], 'Larry')
        files = dict([(f.filename, f) for f in request.POST['files']])
        self.assertEqual(sorted(files), ['file1.txt', 'file2.gif'])

        self.assertEqual(files['file1.txt'].value,
                         b('... contents of file1.txt ...'))
        self.assertEqual(files['file1.txt'].type, 'text/plain')
        self.assertEqual(files['file1.txt'].disposition, 'file')

        self.assertEqual(files['file2.gif'].value,
                         b('...contents of file2.gif...'))
        self.assertEqual(files['file2.gif'].type, 'image/gif')
        self.assertEqual(files['file2.gif'].disposition, 'file')

    def test_upload(self):

        @fiole.post('/upload')
        def post_upload(request):
            return '42'

        environ = {
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'CONTENT_LENGTH': len(FORM_DATA_1),
            'wsgi.input': BytesIO(FORM_DATA_1),
        }
        rv = handle_single_request('GET /upload', **environ)
        self.assertEqual(rv['status'], '405 Method Not Allowed')

        rv = handle_single_request('POST /nothing-here', **environ)
        self.assertEqual(rv['status'], '404 Not Found')

        rv = handle_single_request('POST /upload', **environ)
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('42')])
        self.assertEqual(rv['headers'],
                         [('Content-Type', 'text/html; charset=utf-8'),
                          ('Content-Length', '2')])
        self.assertFalse(rv['errors'])

    def test_cookie(self):

        @fiole.get('/receive')
        def receive_cookies(request):
            cookies = [request.get_cookie(name)
                       for name in sorted(request.cookies)]
            content = cookies
            response = fiole.Response(content, content_type='text/plain',
                                      wrapped=True)
            return response

        @fiole.get('/send')
        def send_cookie(request):
            response = fiole.Response('+ empty +')
            response.set_cookie('nickname', u('gästõn'))
            response.set_cookie('session', 'czpwoe83q8ape2ji23jxnm')
            return response

        gaston = r'g\344st\365n' if PY3 else r'g\303\244st\303\265n'
        # Send cookie
        rv = handle_single_request('GET /send')
        self.assertNoError(rv)
        self.assertEqual(sorted(rv['headers']), [
            ('Content-Length', '9'),
            ('Content-Type', 'text/html; charset=utf-8'),
            ('Set-Cookie', 'nickname="%s"; Path=/' % gaston),
            ('Set-Cookie', 'session=czpwoe83q8ape2ji23jxnm; Path=/'),
        ])
        cookies = [v.split(';')[0]
                   for (k, v) in rv['headers'] if k == 'Set-Cookie']
        cookie = '; '.join(cookies)

        # Receive cookies
        rv = handle_single_request('GET /receive', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']], ['Content-Type'])
        self.assertEqual(rv['data'], ['gästõn', 'czpwoe83q8ape2ji23jxnm'])

    def test_secure_cookie(self):
        fiole.SECRET_KEY = 's e c r e t'

        @fiole.get('/receive')
        def receive_cookies(request):
            cookies = ['%s="%s"' % (name, request.get_cookie(name))
                       for name in request.cookies]
            secure_cookie = request.get_secure_cookie('foo')
            content = cookies + [secure_cookie]
            response = fiole.Response(content, content_type='text/plain',
                                      wrapped=True)
            return response

        @fiole.get('/send_secure')
        def send_secure_cookie(request):
            response = fiole.Response('+ empty +')
            response.set_secure_cookie('foo', u('bär'))
            return response

        @fiole.get('/clear')
        def clear_cookies(request):
            response = fiole.Response('+ empty +')
            for name in request.cookies:
                response.clear_cookie(name)
            return response

        # Send secure cookie
        rv = handle_single_request('GET /send_secure')
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']],
                         ['Content-Type', 'Content-Length', 'Set-Cookie'])
        cookie, expires, path = rv['headers'][2][1].split('; ')
        cookie_kv = cookie.split('|')[0]
        self.assertEqual(cookie_kv, 'foo="YsOkcg==')
        self.assertEqual(expires[:8], 'expires=')
        self.assertEqual(path, 'Path=/')

        # Receive cookies
        rv = handle_single_request('GET /receive', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']], ['Content-Type'])
        (raw_cookie, secure_cookie) = rv['data']
        self.assertEqual(raw_cookie, cookie)
        self.assertEqual(secure_cookie, u('bär'))

        # Clear cookies
        cookie = 'nickname=gaston; ' + cookie
        rv = handle_single_request('GET /clear', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']],
                         ['Content-Type', 'Content-Length',
                          'Set-Cookie', 'Set-Cookie'])
        cookie1, expires1, path1 = rv['headers'][2][1].split('; ')
        cookie2, expires2, path2 = rv['headers'][3][1].split('; ')
        self.assertEqual(sorted([cookie1, cookie2]), ['foo=', 'nickname='])
        self.assertEqual(expires1[:8], 'expires=')
        self.assertEqual(expires2[:8], 'expires=')
        self.assertEqual(path1, 'Path=/')
        self.assertEqual(path2, 'Path=/')

    def test_send_file(self):
        fname = os.path.join(os.path.dirname(__file__), '__init__.py')
        rootdir, fname = os.path.split(fname)

        @fiole.get('/img/logo.png')
        def img_logo(request):
            return fiole.send_file(request, fname, root=rootdir)

        rv = handle_single_request('GET /img/logo.png')
        try:
            self.assertNoError(rv)
            data = b('').join([chunk for chunk in rv['data']])
            with open(os.path.join(rootdir, fname), 'rb') as f:
                self.assertEqual(data, f.read())
        finally:
            rv['data'].close()
<commit_msg>Use the PY3 constant<commit_after># -*- coding: utf-8 -*-
import os.path
import unittest
try:
    from cStringIO import StringIO as BytesIO
except ImportError:
    from io import BytesIO

import fiole
from ._common import (PY3, ENVIRON, handle_single_request,
                      FORM_DATA_CONTENT_TYPE, FORM_DATA_1, FORM_DATA_2)
u = (lambda s: s) if PY3 else (lambda s: s.decode('utf-8'))
b = (lambda s: s.encode('utf-8')) if PY3 else (lambda s: s)

class FioleTestCase(unittest.TestCase):
    _cookie_secret = fiole.SECRET_KEY

    def setUp(self):
        fiole.Fiole.push()
        fiole.SECRET_KEY = self._cookie_secret

    def tearDown(self):
        fiole.Fiole.pop()
        fiole.SECRET_KEY = self._cookie_secret

    def assertNoError(self, response):
        self.assertFalse(response['errors'], msg=response['errors'])
        self.assertEqual(response['status'], '200 OK')

    def test_simple(self):
        @fiole.get('/')
        def index(request):
            return 'Hello World!'
        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('Hello World!')])
        self.assertEqual(rv, {
            'status': '200 OK',
            'headers': [('Content-Type', 'text/html; charset=utf-8'),
                        ('Content-Length', '12')],
            'data': [b('Hello World!')],
            'errors': '',
        })

    def test_request_dispatching(self):
        @fiole.get('/')
        def index(request):
            return request.method

        @fiole.get('/more')
        @fiole.post('/more')
        def more(request):
            return request.method

        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('GET')])

        rv = handle_single_request('POST /')
        self.assertEqual(rv['status'], '405 Method Not Allowed')

        rv = handle_single_request('HEAD /')
        self.assertNoError(rv)
        self.assertFalse(rv['data'])

        rv = handle_single_request('POST /more')
        self.assertEqual(rv['data'], [b('POST')])

        rv = handle_single_request('GET /more')
        self.assertEqual(rv['data'], [b('GET')])

        rv = handle_single_request('DELETE /more')
        self.assertEqual(rv['status'], '405 Method Not Allowed')

    def test_request_route_placeholder(self):
        @fiole.route(u('/<_mot_>/à/<v_42>'))
        def route1(request, _mot_, v_42):
            return u('%s à %s') % (_mot_, v_42)

        rv = handle_single_request(u('GET /table/à/manger'))
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('table à manger')])

    def test_request_route_regex(self):
        @fiole.route(u('/(?P<a>\d+)/plus/(?P<b>\d+)'))
        def addition(request, a, b):
            (a, b) = int(a), int(b)
            return '%d + %d = %d' % (a, b, a + b)

        rv = handle_single_request(u('GET /25/plus/17/'))
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('25 + 17 = 42')])

        rv = handle_single_request(u('GET /navet/plus/tomate/'))
        self.assertEqual(rv['status'], '404 Not Found')

    def test_error_handling(self):
        @fiole.errorhandler(404)
        def not_found(exc):
            return 'not found'

        @fiole.errorhandler(500)
        def internal_server_error(exc):
            return 'internal server error'

        @fiole.get('/')
        def index(request):
            raise fiole.NotFound('Lost')

        @fiole.get('/error')
        def error(request):
            1 // 0

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '404 Not Found')
        self.assertEqual(rv['data'], [b('not found')])
        self.assertFalse(rv['errors'])
        rv = handle_single_request('GET /error')
        self.assertEqual(rv['status'], '500 Internal Server Error')
        self.assertEqual(rv['data'], [b('internal server error')])
        self.assertIn('ZeroDivisionError', rv['errors'])

    def test_user_error_handling(self):
        class MyException(Exception):
            pass

        @fiole.errorhandler(500)
        def system_error(exc):
            if isinstance(exc, MyException):
                return '42'
            return 'internal server error'

        @fiole.get('/')
        def index(request):
            raise MyException()
        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '500 Internal Server Error')
        self.assertEqual(rv['data'], [b('42')])
        self.assertIn('MyException', rv['errors'])

    def test_response(self):

        @fiole.get('/unicode')
        def from_unicode(request):
            return u('Hällo Wörld')

        @fiole.get('/string')
        def from_string(request):
            return b('Hällo Wörld')

        @fiole.get('/args')
        def from_response(request):
            return fiole.Response('Meh', headers=[('X-Foo', 'Testing')],
                                  status=400, content_type='text/plain')

        rv = handle_single_request('GET /unicode')
        self.assertEqual(rv['data'], [b('Hällo Wörld')])
        rv = handle_single_request('GET /string')
        self.assertEqual(rv['data'], [b('Hällo Wörld')])
        rv = handle_single_request('GET /args')
        self.assertEqual(rv, {
            'status': '400 Bad Request',
            'headers': [('X-Foo', 'Testing'),
                        ('Content-Type', 'text/plain; charset=utf-8'),
                        ('Content-Length', '3')],
            'data': [b('Meh')],
            'errors': '',
        })

        rv = fiole.Response('')
        html_utf8 = 'text/html; charset=utf-8'
        self.assertEqual(rv.status, 200)
        self.assertEqual(rv.output, '')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

        rv = fiole.Response('Awesome')
        self.assertEqual(rv.status, 200)
        self.assertEqual(rv.output, 'Awesome')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

        rv = fiole.Response('W00t', status=404)
        self.assertEqual(rv.status, 404)
        self.assertEqual(rv.output, 'W00t')
        self.assertEqual(rv.headers['Content-Type'], html_utf8)

    def test_none_response(self):

        @fiole.get('/')
        def index(request):
            return

        rv = handle_single_request('GET /')
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('')])

    def test_redirect(self):

        @fiole.get('/')
        def index(request):
            raise fiole.Redirect('/far')

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '302 Found')
        self.assertEqual(rv['data'], [b('')])
        self.assertIn(('Location', '/far'), rv['headers'])

    def test_redirect_error_handler(self):

        @fiole.errorhandler(500)
        def handle_500_redirect(request):
            raise fiole.Redirect('/far')

        @fiole.get('/')
        def index(request):
            1 // 0

        rv = handle_single_request('GET /')
        self.assertEqual(rv['status'], '302 Found')
        self.assertEqual(rv['data'], [b('')])
        self.assertIn(('Location', '/far'), rv['headers'])
        # Is it correct?
        self.assertIn('ZeroDivisionError', rv['errors'])

    def test_form_data_1(self):
        environ = dict(ENVIRON)
        environ.update({
            'REQUEST_METHOD': 'POST',
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'CONTENT_LENGTH': len(FORM_DATA_1),
            'wsgi.input': BytesIO(FORM_DATA_1),
        })
        request = fiole.Request(environ)
        self.assertFalse(request.GET)
        self.assertFalse(request.cookies)
        self.assertEqual(sorted(request.POST), ['files', 'submit-name'])
        self.assertEqual(request.POST['submit-name'], 'Larry')
        self.assertEqual(request.POST['files'].filename, 'file1.txt')
        self.assertEqual(request.POST['files'].file.read(),
                         b('... contents of file1.txt ...'))
        self.assertEqual(request.POST['files'].value,
                         b('... contents of file1.txt ...'))
        self.assertEqual(request.POST['files'].type, 'text/plain')
        self.assertEqual(request.POST['files'].disposition, 'form-data')

    def test_form_data_2(self):
        environ = dict(ENVIRON)
        environ.update({
            'REQUEST_METHOD': 'POST',
            'CONTENT_LENGTH': len(FORM_DATA_2),
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'wsgi.input': BytesIO(FORM_DATA_2),
        })
        request = fiole.Request(environ)
        self.assertFalse(request.GET)
        self.assertFalse(request.cookies)
        self.assertEqual(sorted(request.POST), ['files', 'submit-name'])
        self.assertEqual(request.POST['submit-name'], 'Larry')
        files = dict([(f.filename, f) for f in request.POST['files']])
        self.assertEqual(sorted(files), ['file1.txt', 'file2.gif'])

        self.assertEqual(files['file1.txt'].value,
                         b('... contents of file1.txt ...'))
        self.assertEqual(files['file1.txt'].type, 'text/plain')
        self.assertEqual(files['file1.txt'].disposition, 'file')

        self.assertEqual(files['file2.gif'].value,
                         b('...contents of file2.gif...'))
        self.assertEqual(files['file2.gif'].type, 'image/gif')
        self.assertEqual(files['file2.gif'].disposition, 'file')

    def test_upload(self):

        @fiole.post('/upload')
        def post_upload(request):
            return '42'

        environ = {
            'CONTENT_TYPE': FORM_DATA_CONTENT_TYPE,
            'CONTENT_LENGTH': len(FORM_DATA_1),
            'wsgi.input': BytesIO(FORM_DATA_1),
        }
        rv = handle_single_request('GET /upload', **environ)
        self.assertEqual(rv['status'], '405 Method Not Allowed')

        rv = handle_single_request('POST /nothing-here', **environ)
        self.assertEqual(rv['status'], '404 Not Found')

        rv = handle_single_request('POST /upload', **environ)
        self.assertNoError(rv)
        self.assertEqual(rv['data'], [b('42')])
        self.assertEqual(rv['headers'],
                         [('Content-Type', 'text/html; charset=utf-8'),
                          ('Content-Length', '2')])
        self.assertFalse(rv['errors'])

    def test_cookie(self):

        @fiole.get('/receive')
        def receive_cookies(request):
            cookies = [request.get_cookie(name)
                       for name in sorted(request.cookies)]
            content = cookies
            response = fiole.Response(content, content_type='text/plain',
                                      wrapped=True)
            return response

        @fiole.get('/send')
        def send_cookie(request):
            response = fiole.Response('+ empty +')
            response.set_cookie('nickname', u('gästõn'))
            response.set_cookie('session', 'czpwoe83q8ape2ji23jxnm')
            return response

        gaston = r'g\344st\365n' if PY3 else r'g\303\244st\303\265n'
        # Send cookie
        rv = handle_single_request('GET /send')
        self.assertNoError(rv)
        self.assertEqual(sorted(rv['headers']), [
            ('Content-Length', '9'),
            ('Content-Type', 'text/html; charset=utf-8'),
            ('Set-Cookie', 'nickname="%s"; Path=/' % gaston),
            ('Set-Cookie', 'session=czpwoe83q8ape2ji23jxnm; Path=/'),
        ])
        cookies = [v.split(';')[0]
                   for (k, v) in rv['headers'] if k == 'Set-Cookie']
        cookie = '; '.join(cookies)

        # Receive cookies
        rv = handle_single_request('GET /receive', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']], ['Content-Type'])
        self.assertEqual(rv['data'], ['gästõn', 'czpwoe83q8ape2ji23jxnm'])

    def test_secure_cookie(self):
        fiole.SECRET_KEY = 's e c r e t'

        @fiole.get('/receive')
        def receive_cookies(request):
            cookies = ['%s="%s"' % (name, request.get_cookie(name))
                       for name in request.cookies]
            secure_cookie = request.get_secure_cookie('foo')
            content = cookies + [secure_cookie]
            response = fiole.Response(content, content_type='text/plain',
                                      wrapped=True)
            return response

        @fiole.get('/send_secure')
        def send_secure_cookie(request):
            response = fiole.Response('+ empty +')
            response.set_secure_cookie('foo', u('bär'))
            return response

        @fiole.get('/clear')
        def clear_cookies(request):
            response = fiole.Response('+ empty +')
            for name in request.cookies:
                response.clear_cookie(name)
            return response

        # Send secure cookie
        rv = handle_single_request('GET /send_secure')
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']],
                         ['Content-Type', 'Content-Length', 'Set-Cookie'])
        cookie, expires, path = rv['headers'][2][1].split('; ')
        cookie_kv = cookie.split('|')[0]
        self.assertEqual(cookie_kv, 'foo="YsOkcg==')
        self.assertEqual(expires[:8], 'expires=')
        self.assertEqual(path, 'Path=/')

        # Receive cookies
        rv = handle_single_request('GET /receive', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']], ['Content-Type'])
        (raw_cookie, secure_cookie) = rv['data']
        self.assertEqual(raw_cookie, cookie)
        self.assertEqual(secure_cookie, u('bär'))

        # Clear cookies
        cookie = 'nickname=gaston; ' + cookie
        rv = handle_single_request('GET /clear', HTTP_COOKIE=cookie)
        self.assertNoError(rv)
        self.assertEqual([h for (h, v) in rv['headers']],
                         ['Content-Type', 'Content-Length',
                          'Set-Cookie', 'Set-Cookie'])
        cookie1, expires1, path1 = rv['headers'][2][1].split('; ')
        cookie2, expires2, path2 = rv['headers'][3][1].split('; ')
        self.assertEqual(sorted([cookie1, cookie2]), ['foo=', 'nickname='])
        self.assertEqual(expires1[:8], 'expires=')
        self.assertEqual(expires2[:8], 'expires=')
        self.assertEqual(path1, 'Path=/')
        self.assertEqual(path2, 'Path=/')

    def test_send_file(self):
        fname = os.path.join(os.path.dirname(__file__), '__init__.py')
        rootdir, fname = os.path.split(fname)

        @fiole.get('/img/logo.png')
        def img_logo(request):
            return fiole.send_file(request, fname, root=rootdir)

        rv = handle_single_request('GET /img/logo.png')
        try:
            self.assertNoError(rv)
            data = b('').join([chunk for chunk in rv['data']])
            with open(os.path.join(rootdir, fname), 'rb') as f:
                self.assertEqual(data, f.read())
        finally:
            rv['data'].close()
<|endoftext|>