~didrocks/unity/altf10

from BaseHTTPServer import (
    HTTPServer,
    BaseHTTPRequestHandler,
    )
import httplib
import logging
import os
from signal import SIGKILL
import simplejson
from StringIO import StringIO
import sys
from unittest import TestCase
from urlparse import urlparse
from urlparse import parse_qs

from testtools import ExpectedException

from grackle.client import (
    GrackleClient,
    UnsupportedOrder,
    )


def threaded_messages(messages):
    threads = {}
    count = 0
    pending = []
    for message in messages:
        if message.get('in_reply_to') is None:
            threads[message['message_id']] = [message]
            count += 1
        else:
            pending.append(message)
    for message in pending:
        threads[message['in_reply_to']].append(message)
    return threads.values()


class GrackleStore:
    """A memory-backed message store."""

    def __init__(self, messages):
        """Constructor."""
        self.messages = messages

    def get_messages(self, archive_id, query_string):
        """Return matching messages.

        :param archive_id: The archive to retrieve from.
        :param query_string: Contains 'parameters', which is a JSON-format
            string describing parameters.
        """
        query = parse_qs(query_string)
        parameters = simplejson.loads(query['parameters'][0])
        order = parameters.get('order')
        messages = self.messages[archive_id]
        if order is not None:
            if order not in SUPPORTED_ORDERS:
                raise UnsupportedOrder
            elif order.startswith('thread_'):
                threaded = threaded_messages(messages)
                messages = []
                if order == 'thread_subject':
                    threaded.sort(key=lambda t: t[0]['subject'])
                if order == 'thread_oldest':
                    threaded.sort(key=lambda t: min(m['date'] for m in t))
                if order == 'thread_newest':
                    threaded.sort(key=lambda t: max(m['date'] for m in t))
                for thread in threaded:
                    messages.extend(thread)
            else:
                messages.sort(key=lambda m: m[order])
        new_messages = []
        for message in messages:
            if (not parameters['include_hidden']
                and message.get('hidden', False)):
                continue

            if ('message_ids' in parameters
                and message['message_id'] not in parameters['message_ids']):
                continue
            message = dict(message)
            if 'headers' in parameters:
                headers = dict(
                    (k, v) for k, v in message['headers'].iteritems()
                    if k in parameters['headers'])
                message['headers'] = headers
            max_body = parameters.get('max_body_length')
            if max_body is not None:
                message['body'] = message['body'][:max_body]
            new_messages.append(message)
        messages = new_messages
        limit = parameters.get('limit', 100)
        memo = parameters.get('memo')
        message_id_indices = dict(
            (m['message_id'], idx) for idx, m in enumerate(messages))
        if memo is None:
            start = 0
        else:
            start = message_id_indices[memo.encode('rot13')]
        if start > 0:
            previous_memo = messages[start - 1]['message_id'].encode('rot13')
        else:
            previous_memo = None
        end = min(start + limit, len(messages))
        if end < len(messages):
            next_memo = messages[end]['message_id'].encode('rot13')
        else:
            next_memo = None
        messages = messages[start:end]

        response = {
            'messages': messages,
            'next_memo': next_memo,
            'previous_memo': previous_memo
            }
        return response


class ForkedFake:
    """A Grackle service fake, as a ContextManager."""

    def __init__(self, port, messages=None, write_logs=False):
        """Constructor.
        :param port: The tcp port to use
        :param messages: A dict of lists of dicts representing messages.  The
            outer dict represents the archive, the list represents the list of
            messages for that archive.
        :param write_logs: If true, log messages will be written to stdout.
        """
        self.pid = None
        self.port = port
        if messages is None:
            self.messages = {}
        else:
            self.messages = messages
        self.read_end, self.write_end = os.pipe()
        self.write_logs = write_logs

    @staticmethod
    def from_client(client, messages=None):
        """Instantiate a ForkedFake from the client.

        :param port: The client  to provide service for.
        :param messages: A dict of lists of dicts representing messages.  The
            outer dict represents the archive, the list represents the list of
            messages for that archive.
        """
        return ForkedFake(client.port, messages)

    def is_ready(self):
        """Tell the parent process that the server is ready for writes."""
        os.write(self.write_end, 'asdf')

    def __enter__(self):
        """Run the service.

        Fork and start a server in the child.  Return when the server is ready
        for use."""
        pid = os.fork()
        if pid == 0:
            self.start_server()
        self.pid = pid
        os.read(self.read_end, 1)
        return

    def start_server(self):
        """Start the HTTP server."""
        service = HTTPServer(('', self.port), FakeGrackleRequestHandler)
        service.store = GrackleStore(self.messages)
        for archive_id, messages in service.store.messages.iteritems():
            for message in messages:
                message.setdefault('headers', {})
        self.is_ready()
        if self.write_logs:
            logging.basicConfig(
                stream=sys.stderr, level=logging.INFO)
        service.serve_forever()

    def __exit__(self, exc_type, exc_val, traceback):
        os.kill(self.pid, SIGKILL)


SUPPORTED_ORDERS = set(
    ['date', 'author', 'subject', 'thread_newest', 'thread_oldest',
     'thread_subject'])


class FakeGrackleRequestHandler(BaseHTTPRequestHandler):
    """A request handler that forwards to server.store."""

    def __init__(self, *args, **kwargs):
        """Constructor.  Sets up logging."""
        self.logger = logging.getLogger('http')
        BaseHTTPRequestHandler.__init__(self, *args, **kwargs)

    def do_POST(self):
        """Create a message on POST."""
        message = self.rfile.read(int(self.headers['content-length']))
        if message == 'This is a message':
            self.send_response(httplib.CREATED)
            self.end_headers()
            self.wfile.close()
        else:
            self.send_error(httplib.BAD_REQUEST)

    def do_GET(self):
        """Retrieve a list of messages on GET."""
        scheme, netloc, path, params, query_string, fragments = (
            urlparse(self.path))
        parts = path.split('/')
        if parts[1] == 'archive':
            try:
                response = self.server.store.get_messages(
                    parts[2], query_string)
                self.send_response(httplib.OK)
                self.end_headers()
                self.wfile.write(simplejson.dumps(response))
            except UnsupportedOrder:
                self.send_response(httplib.BAD_REQUEST)
                self.wfile.write('Unsupported order')
                return

    def log_message(self, format, *args):
        """Override log_message to use standard Python logging."""
        message = "%s - - [%s] %s\n" % (
            self.address_string(), self.log_date_time_string(), format % args)
        self.logger.info(message)


class TestPutMessage(TestCase):

    def test_put_message(self):
        client = GrackleClient('localhost', 8436)
        with ForkedFake.from_client(client):
            client.put_message('arch1', 'asdf', StringIO('This is a message'))
            with ExpectedException(Exception, 'wtf'):
                client.put_message('arch1', 'asdf',
                    StringIO('This is not a message'))


class TestGetMessages(TestCase):

    def assertIDOrder(self, ids, messages):
        self.assertEqual(ids, [m['message_id'] for m in messages])

    def assertMessageIDs(self, ids, messages):
        self.assertIDOrder(
            sorted(ids), sorted(messages, key=lambda m: m['message_id']))

    def test_get_messages(self):
        client = GrackleClient('localhost', 8435)
        with ForkedFake.from_client(client,
            {'baz':
            [{'message_id': 'foo'},
             {'message_id': 'bar'}]}):
            response = client.get_messages('baz')
        self.assertEqual(['bar', 'foo'], sorted(m['message_id'] for m in
            response['messages']))
        self.assertIs(None, response['next_memo'])
        self.assertIs(None, response['previous_memo'])

    def test_get_messages_by_id(self):
        client = GrackleClient('localhost', 8437)
        with ForkedFake.from_client(client,
            {'baz':
            [{'message_id': 'foo'},
             {'message_id': 'bar'}]}):
            response = client.get_messages('baz', message_ids=['foo'])
        message, = response['messages']
        self.assertEqual('foo', message['message_id'])

    def test_get_messages_batching(self):
        client = GrackleClient('localhost', 8438)
        with ForkedFake.from_client(client,
            {'baz':
            [{'message_id': 'foo'},
             {'message_id': 'bar'}]}):
            response = client.get_messages('baz', limit=1)
            self.assertEqual(1, len(response['messages']))
            messages = response['messages']
            response = client.get_messages(
                'baz', limit=1, memo=response['next_memo'])
            self.assertEqual(1, len(response['messages']))
            messages.extend(response['messages'])
            self.assertMessageIDs(['foo', 'bar'], messages)

    def get_messages_member_order_test(self, key):
        client = GrackleClient('localhost', 8439)
        with ForkedFake.from_client(client,
                {'baz': [{'message_id': 'foo', key: '2011-03-25'},
                 {'message_id': 'bar', key: '2011-03-24'}]}):
            response = client.get_messages('baz')
            self.assertIDOrder(['foo', 'bar'], response['messages'])
            response = client.get_messages('baz', order=key)
            self.assertIDOrder(['bar', 'foo'], response['messages'])

    def test_get_messages_date_order(self):
        self.get_messages_member_order_test('date')

    def test_get_messages_author_order(self):
        self.get_messages_member_order_test('author')

    def test_get_messages_subject_order(self):
        self.get_messages_member_order_test('subject')

    def test_get_messages_thread_subject_order(self):
        client = GrackleClient('localhost', 8439)
        with ForkedFake.from_client(client, {'baz': [
            {'message_id': 'bar', 'subject': 'y'},
            {'message_id': 'qux', 'subject': 'z'},
            {'message_id': 'foo', 'subject': 'x', 'in_reply_to': 'qux'},
            ]}):
            response = client.get_messages('baz')
            self.assertIDOrder(['bar', 'qux', 'foo'], response['messages'])
            response = client.get_messages('baz', order='subject')
            self.assertIDOrder(['foo', 'bar', 'qux'], response['messages'])
            response = client.get_messages('baz', order='thread_subject')
            self.assertIDOrder(['bar', 'qux', 'foo'], response['messages'])

    def test_get_messages_thread_oldest_order(self):
        client = GrackleClient('localhost', 8439)
        with ForkedFake.from_client(client, {'baz': [
            {'message_id': 'bar', 'date': 'x'},
            {'message_id': 'qux', 'date': 'z'},
            {'message_id': 'foo', 'date': 'y', 'in_reply_to': 'qux'},
            ]}):
            response = client.get_messages('baz')
            self.assertIDOrder(['bar', 'qux', 'foo'], response['messages'])
            response = client.get_messages('baz', order='date')
            self.assertIDOrder(['bar', 'foo', 'qux'], response['messages'])
            response = client.get_messages('baz', order='thread_oldest')
            self.assertIDOrder(['bar', 'qux', 'foo'], response['messages'])

    def test_get_messages_thread_newest_order(self):
        client = GrackleClient('localhost', 8439)
        with ForkedFake.from_client(client, {'baz': [
            {'message_id': 'bar', 'date': 'x'},
            {'message_id': 'qux', 'date': 'w'},
            {'message_id': 'foo', 'date': 'y', 'in_reply_to': 'bar'},
            {'message_id': 'baz', 'date': 'z', 'in_reply_to': 'qux'},
            ]}):
            response = client.get_messages('baz', order='date')
            self.assertIDOrder(
                ['qux', 'bar', 'foo', 'baz'], response['messages'])
            response = client.get_messages('baz', order='thread_newest')
            self.assertIDOrder(
                ['bar', 'foo', 'qux', 'baz'], response['messages'])

    def test_get_messages_unsupported_order(self):
        client = GrackleClient('localhost', 8439)
        with ForkedFake.from_client(client,
                {'baz': [{'message_id': 'foo', 'date': '2011-03-25'},
                 {'message_id': 'bar', 'date': '2011-03-24'}]}):
            with ExpectedException(UnsupportedOrder, ''):
                client.get_messages('baz', order='nonsense')

    def test_get_messages_headers_no_headers(self):
        client = GrackleClient('localhost', 8440)
        with ForkedFake.from_client(client,
            {'baz': [
                {'message_id': 'foo'}
            ]}):
            response = client.get_messages('baz', headers=[
                'Subject', 'Date', 'X-Launchpad-Message-Rationale'])
        first_message = response['messages'][0]
        self.assertEqual('foo', first_message['message_id'])
        self.assertEqual({}, first_message['headers'])

    def test_get_messages_headers_exclude_headers(self):
        client = GrackleClient('localhost', 8441)
        with ForkedFake.from_client(client,
            {'baz': [
                {'message_id': 'foo', 'headers': {'From': 'me'}}
            ]}):
            response = client.get_messages('baz', headers=[
                'Subject', 'Date', 'X-Launchpad-Message-Rationale'])
        first_message = response['messages'][0]
        self.assertEqual('foo', first_message['message_id'])
        self.assertEqual({}, first_message['headers'])

    def test_get_messages_headers_include_headers(self):
        client = GrackleClient('localhost', 8442)
        with ForkedFake.from_client(client,
            {'baz': [
                {'message_id': 'foo', 'headers': {'From': 'me', 'To': 'you'}}
            ]}):
            response = client.get_messages('baz', headers=[
                'From', 'To'])
        first_message = response['messages'][0]
        self.assertEqual('foo', first_message['message_id'])
        self.assertEqual({'From': 'me', 'To': 'you'}, first_message['headers'])

    def test_get_messages_max_body_length(self):
        client = GrackleClient('localhost', 8443)
        with ForkedFake.from_client(client,
            {'baz': [
                {'message_id': 'foo', 'body': u'abcdefghi'}
            ]}):
            response = client.get_messages('baz', max_body_length=3)
        first_message = response['messages'][0]
        self.assertEqual('abc', first_message['body'])

    def test_include_hidden(self):
        client = GrackleClient('localhost', 8444)
        with ForkedFake.from_client(client,
            {'baz': [
                {'message_id': 'foo', 'hidden': True},
                {'message_id': 'bar', 'hidden': False}
            ]}):
            response = client.get_messages('baz', include_hidden=True)
            self.assertMessageIDs(['bar', 'foo'], response['messages'])
            response = client.get_messages('baz', include_hidden=False)
            self.assertMessageIDs(['bar'], response['messages'])