~didrocks/unity/altf10

from BaseHTTPServer import (

    HTTPServer,

    BaseHTTPRequestHandler,

    )

import httplib

import os

from signal import SIGKILL

from StringIO import StringIO

from unittest import TestCase

from testtools import ExpectedException

from grackle.client import GrackleClient

class Forked:

    def __init__(self, func_or_method, *args):

        self.func_or_method = func_or_method

        self.pid = None

        self.args = args

    def __enter__(self):

        pid = os.fork()

        if pid != 0:

            self.pid = pid

            return

        self.func_or_method(*self.args)

    def __exit__(self, exc_type, exc_val, traceback):

        os.kill(self.pid, SIGKILL)

class FakeGrackleRequestHandler(BaseHTTPRequestHandler):

    def do_PUT(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 run_service(port):

    service = HTTPServer(('', port), FakeGrackleRequestHandler)

    service.serve_forever()

class TestPutMessage(TestCase):

    def test_put_message(self):

        client = GrackleClient('localhost', 8435)

        with Forked(run_service, client.port):

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