~didrocks/unity/altf10

from BaseHTTPServer import (

    HTTPServer,

    BaseHTTPRequestHandler,

    )

import httplib

import os

from signal import SIGKILL

import simplejson

from StringIO import StringIO

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:

    def __init__(self, messages):

        self.messages = messages

    def get_messages(self, archive_id, query_string):

        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])

        messages = [m for m in messages

                    if 'message_ids' not in parameters or

                    m['message_id'] in parameters['message_ids']]

        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]

        new_messages = []

        for message in messages:

            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)

        response = {

            'messages': new_messages,

            'next_memo': next_memo,

            'previous_memo': previous_memo

            }

        return response

class ForkedFake:

    def __init__(self, port, messages=None):

        self.pid = None

        self.port = port

        self.messages = messages

        self.read_end, self.write_end = os.pipe()

    def is_ready(self):

        os.write(self.write_end, 'asdf')

    def __enter__(self):

        pid = os.fork()

        if pid == 0:

            self.start_server()

        self.pid = pid

        os.read(self.read_end, 1)

        return

    def start_server(self):

        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()

        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):

    def do_POST(self):

        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):

        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 fake_grackle_service(client, messages=None):

    if messages is None:

        messages = {}

    return ForkedFake(client.port, messages)

class TestPutMessage(TestCase):

    def test_put_message(self):

        client = GrackleClient('localhost', 8436)

        with fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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 fake_grackle_service(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', 8440)

        with fake_grackle_service(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', 8440)

        with fake_grackle_service(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', 8440)

        with fake_grackle_service(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'])