~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 datetime

from email.utils import formatdate

import os

from tempfile import _RandomNameSequence

import time

import unittest

from dateutil.tz import (

    tzoffset,

    tzutc,

    )

from pycassa.util import convert_uuid_to_time

from grackle.model import (

    CassandraConnection,

    _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 = time.mktime(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['content'])

        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['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)

        self.assertMessages(

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

    def test_limit(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 4)

        self.assertMessages(

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

    def test_order(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 4)

        self.assertMessages(

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

    def test_batching_forward(self):

        conn, archive = self.makeArchive()

        self.makeMessages(conn, archive, 5)

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

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

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

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

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

        self.assertMessages([4], messages)

        self.assertIs(None, start)