~unity-2d-team/unity-2d/Shell-MultiMonitor

# Copyright (c) 2012 Canonical Ltd

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

#

# You should have received a copy of the GNU Affero General Public

# License along with this program. If not, see

# <http://www.gnu.org/licenses/>.

__metaclass__ = type

import calendar

import datetime

from email.utils import formatdate

import os

from tempfile import _RandomNameSequence

import unittest

from dateutil.tz import (

    tzoffset,

    tzutc,

    )

from pycassa.util import convert_uuid_to_time

from grackle.model import (

    CassandraConnection,

    _format_message,

    _parse_message,

    )

TEMPLATE_MESSAGE = """\

From: sysadmin@example.com

To: developer@example.com

Subject: Everything is broken

Date: {date}

Message-Id: {id}

Help, everything has just broken.

"""

TEST_MESSAGE = TEMPLATE_MESSAGE.format(

    date='Sat, 1 Jan 2000 11:02:34 +1100',

    id='<aaaaaaaaaaaaa@example.com>')

class TestParseMessage(unittest.TestCase):

    def test_works(self):

        # _parse_message extracts interesting fields. It also parses the

        # date and returns it separately.

        date, msg = _parse_message(TEST_MESSAGE)

        self.assertEqual('sysadmin@example.com', msg['from'])

        self.assertEqual('developer@example.com', msg['to'])

        self.assertEqual('Everything is broken', msg['subject'])

        self.assertEqual('2000-01-01T11:02:34+11:00', msg['date'])

        self.assertEqual('<aaaaaaaaaaaaa@example.com>', msg['message-id'])

        self.assertEqual(

            datetime.datetime(

                2000, 1, 1, 11, 2, 34, tzinfo=tzoffset('', 39600)),

            date)

class TestAddMessage(unittest.TestCase):

    def test_add_message(self):

        c = CassandraConnection(

            os.environ['GRACKLE_TEST_KEYSPACE'], ['localhost:9160'])

        archive = next(_RandomNameSequence())

        # Write the message out to Cassandra, and read it back in.

        key = c.add_message(archive, TEST_MESSAGE)

        cmsg = c.messages.get(key)

        # The archive should contain a single message, a reference to

        # our new key.

        archive_messages = c.archive_messages.get(archive).items()

        self.assertEqual(1, len(archive_messages))

        self.assertEqual(key, archive_messages[0][1])

        # The key in archive_message is a TimeUUID for the Date field in

        # the message.

        utctime = calendar.timegm(datetime.datetime(

            2000, 1, 1, 0, 2, 34, tzinfo=tzutc()).timetuple())

        self.assertEqual(

            utctime,

            convert_uuid_to_time(archive_messages[0][0]))

        # The stored message contains the full original text of the

        # message, as well as interesting fields parsed out.

        self.assertEqual(TEST_MESSAGE, cmsg['raw'])

        parsed_message = _parse_message(TEST_MESSAGE)[1]

        for key, value in parsed_message.iteritems():

            self.assertEqual(value, cmsg[key])

class TestGetMessages(unittest.TestCase):

    def assertMessages(self, expected_ids, messages):

        expected_msgids = [

            '<message%d@example.com>' % id for id in expected_ids]

        actual_msgids = [msg['headers']['message-id'] for msg in messages]

        self.assertEqual(expected_msgids, actual_msgids)

    def makeMessages(self, conn, archive, count):

        return [

            conn.add_message(

                archive,

                TEMPLATE_MESSAGE.format(

                    date=formatdate(i * 100),

                    id='<message%d@example.com>' % i))

            for i in range(count)]

    def makeArchive(self):

        conn = CassandraConnection(

            os.environ['GRACKLE_TEST_KEYSPACE'], ['localhost:9160'])

        archive = next(_RandomNameSequence())

        return conn, archive

    def test_single_message(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 1)

        # We get a single message when we ask for it.

        messages = conn.get_messages(archive, 'date', 1, '')[1]

        self.assertMessages([0], messages)

        cmsg = messages[0]

        # The raw message matches what we generated, and the result is

        # formatted according to the default settings.

        expected_raw = TEMPLATE_MESSAGE.format(

            date=formatdate(0), id='<message0@example.com>')

        self.assertEqual(expected_raw, messages[0]['raw'])

        pmsg = _parse_message(expected_raw)[1]

        pmsg['raw'] = expected_raw

        self.assertEqual(

            _format_message(

                pmsg,

                headers=['date', 'from', 'subject', 'message-id'],

                include_raw=True),

            cmsg)

    def test_limit(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 4)

        self.assertMessages(

            [0, 1], conn.get_messages(archive, 'date', 2, '')[1])

    def test_order(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 4)

        self.assertMessages(

            [3, 2], conn.get_messages(archive, '-date', 2, '')[1])

    def test_batching_forward(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 5)

        prev, messages, next = conn.get_messages(archive, 'date', 2, '')

        self.assertMessages([0, 1], messages)

        prev, messages, next = conn.get_messages(archive, 'date', 2, next)

        self.assertMessages([2, 3], messages)

        prev, messages, next = conn.get_messages(archive, 'date', 2, next)

        self.assertMessages([4], messages)

        prev, messages, next = conn.get_messages(archive, 'date', 2, next)

        self.assertIs(None, prev)

        self.assertMessages([], messages)

        self.assertIs(None, next)

    def test_batching_backward(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 5)

        prev, messages, next = conn.get_messages(archive, 'date', 2, '')

        self.assertMessages([0, 1], messages)

        prev, messages, next = conn.get_messages(archive, 'date', 2, next)

        self.assertMessages([2, 3], messages)

        prev, messages, next = conn.get_messages(

            archive, 'date', 2, prev, backward=True)

        self.assertMessages([0, 1], messages)

        prev, messages, next = conn.get_messages(

            archive, 'date', 2, prev, backward=True)

        self.assertIs(None, prev)

        self.assertMessages([], messages)

        self.assertIs(None, next)

    def test_date_filter(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 10)

        start = datetime.datetime.utcfromtimestamp(250).replace(

            tzinfo=tzutc())

        finish = datetime.datetime.utcfromtimestamp(500).replace(

            tzinfo=tzutc())

        prev, messages, next = conn.get_messages(

            archive, 'date', 2, '', start_date=start, finish_date=finish)

        self.assertMessages([3, 4], messages)

        prev, messages, next = conn.get_messages(

            archive, 'date', 2, next, start_date=start, finish_date=finish)

        self.assertMessages([5], messages)

class TestMessageFormatter(unittest.TestCase):

    def test_all(self):

        parsed = _parse_message(TEST_MESSAGE)[1]

        parsed['raw'] = TEST_MESSAGE

        formatted = _format_message(

            parsed,

            headers=['date', 'from', 'subject', 'message-id'],

            include_raw=True)

        expected = {

            'headers': {

                'date': '2000-01-01T11:02:34+11:00',

                'from': 'sysadmin@example.com',

                'message-id': '<aaaaaaaaaaaaa@example.com>',

                'subject': 'Everything is broken',

                },

            'raw': TEST_MESSAGE,

            }

        self.assertEqual(expected, formatted)