'''Tests the websocket middleware in pulsar.apps.ws.'''
from random import randint
import struct

from pulsar import send
from pulsar.apps.ws import Frame, WebSocket, int2bytes, i2b,\
                            WebSocketProtocolError, FrameParser, HttpClient
from pulsar.apps.test import unittest, dont_run_with_thread

from .manage import server


class TestWebSocketThread(unittest.TestCase):
    app = None
    concurrency = 'thread'
    
    @classmethod
    def setUpClass(cls):
        s = server(bind='127.0.0.1:0', name=cls.__name__,
                   concurrency=cls.concurrency)
        outcome = send('arbiter', 'run', s)
        yield outcome
        cls.app = outcome.result
        cls.uri = 'http://{0}:{1}'.format(*cls.app.address)
        cls.ws_uri = 'ws://{0}:{1}/data'.format(*cls.app.address)
        cls.ws_echo = 'ws://{0}:{1}/echo'.format(*cls.app.address)
        
    @classmethod
    def tearDownClass(cls):
        if cls.app is not None:
            yield send('arbiter', 'kill_actor', cls.app.name)
    
    def testHyBiKey(self):
        w = WebSocket(None)
        v = w.challenge_response('dGhlIHNhbXBsZSBub25jZQ==')
        self.assertEqual(v, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
        
    def testBadRequests(self):
        c = HttpClient()
        outcome = c.post(self.ws_uri)
        yield outcome
        response = outcome.result
        self.assertEqual(response.status_code, 400)
        #
        outcome = c.get(self.ws_uri, headers=[('Sec-Websocket-Key', '')])
        yield outcome
        response = outcome.result
        self.assertEqual(response.status_code, 400)
        #
        outcome = c.get(self.ws_uri, headers=[('Sec-Websocket-Key', 'bla')])
        yield outcome
        response = outcome.result
        self.assertEqual(response.status_code, 400)
        #
        outcome = c.get(self.ws_uri, headers=[('Sec-Websocket-version', 'xxx')])
        yield outcome
        response = outcome.result
        self.assertEqual(response.status_code, 400)
    
    def testUpgrade(self):
        c = HttpClient()
        outcome = c.get(self.ws_echo)
        yield outcome
        ws = outcome.result
        response = ws.handshake 
        self.assertEqual(response.status_code, 101)
        self.assertEqual(response.headers['upgrade'], 'websocket')
        # Send a message to the websocket
        outcome = ws.execute('Hi there!')
        yield outcome
        response = outcome.result
        self.assertEqual(response.body, 'Hi there!')
        self.assertFalse(response.masking_key)
        

@dont_run_with_thread
class TestWebSocketProcess(TestWebSocketThread):
    concurrency = 'process'
    
    
class FrameTest(unittest.TestCase):
    
    @classmethod
    def setUpClass(cls):
        cls.bdata = b''.join((i2b(randint(0,255)) for v in range(256)))
        SIZE = 64*1024
        cls.large_bdata = b''.join((i2b(randint(0,255)) for v in range(SIZE)))
        
    def testBadFrame(self):
        self.assertRaises(WebSocketProtocolError, Frame, version=809)
        self.assertRaises(WebSocketProtocolError, Frame, opcode=56)
        self.assertRaises(WebSocketProtocolError, Frame)
        self.assertRaises(WebSocketProtocolError, Frame, 'bla',
                          masking_key='foo')
        self.assertRaises(WebSocketProtocolError, Frame, 'bla',
                          masking_key='fooox')
        
    def testControlFrames(self):
        f = Frame.close()
        self.assertEqual(f.opcode, 0x8)
        self.assertTrue(f.payload_length <= 125)
        f = Frame.ping('Hello')
        self.assertEqual(f.opcode, 0x9)
        self.assertEqual(int2bytes(0x89,0x05,0x48,0x65,0x6c,0x6c,0x6f), f.msg)
        self.assertTrue(f.payload_length <= 125)
        r = f.on_received()
        self.assertTrue(r)
        self.assertEqual(r.opcode, 0xA)
        f = Frame.pong()
        self.assertEqual(f.opcode, 0xA)
        self.assertTrue(f.payload_length <= 125)
        
    def testUnmaskedDataFrame(self):
        f = Frame('Hello', final=True)
        self.assertEqual(f.opcode, 0x1)
        self.assertEqual(f.payload_length, 5)
        self.assertFalse(f.masked)
        self.assertEqual(len(f.msg), 7)
        self.assertEqual(int2bytes(0x81,0x05,0x48,0x65,0x6c,0x6c,0x6f), f.msg)
        f1 = Frame('Hel')
        f2 = Frame.continuation('lo', final=True)
        self.assertEqual(int2bytes(0x01,0x03,0x48,0x65,0x6c), f1.msg)
        self.assertEqual(int2bytes(0x80,0x02,0x6c,0x6f), f2.msg)
        
    def testBinaryDataFrame(self):
        f = Frame(self.bdata, opcode=0x2, final=True)
        self.assertEqual(f.opcode, 0x2)
        self.assertEqual(f.payload_length, 256)
        self.assertFalse(f.masked)
        self.assertEqual(len(f.msg), 260)
        self.assertEqual(struct.pack('!BBH',0x82,0x7E,0x0100), f.msg[:4])
        f = Frame(self.large_bdata, opcode=0x2, final=True)
        self.assertEqual(f.opcode, 0x2)
        self.assertEqual(f.payload_length, len(self.large_bdata))
        self.assertFalse(f.masked)
        self.assertEqual(len(f.msg), len(self.large_bdata)+10)
        self.assertEqual(struct.pack('!BBQ',0x82,0x7F,0x0000000000010000),
                         f.msg[:10])
            
    def testMaskData(self):
        masking_key = int2bytes(0x37,0xfa,0x21,0x3d)
        f = Frame('Hello', masking_key=masking_key, final=True)
        self.assertTrue(f.final)
        msg = int2bytes(0x81,0x85,0x37,0xfa,0x21,0x3d,0x7f,0x9f,0x4d,0x51,0x58)
        self.assertTrue(f.masked)
        self.assertEqual(msg, f.msg)
        
    def testParser(self):
        p = FrameParser()
        self.assertEqual(p.execute(b''), None)
        frame = Frame('Hello', final=True)
        self.assertRaises(WebSocketProtocolError, p.execute, frame.msg)
        frame = Frame('Hello', masking_key='ciao', final=True)
        pframe = p.execute(frame.msg)
        self.assertTrue(pframe)
        self.assertEqual(pframe.body, 'Hello')
        self.assertEqual(pframe.payload_length, 5)
        self.assertEqual(pframe.masking_key, b'ciao')
    
    def testParserBinary(self):
        p = FrameParser()
        frame = Frame(self.bdata, opcode=0x2, final=True, masking_key='ciao')
        pframe = p.execute(frame.msg)
        self.assertTrue(pframe)
        self.assertEqual(pframe.payload_length, 256)
        self.assertEqual(pframe.body, self.bdata)
        frame = Frame(self.large_bdata, opcode=0x2, final=True,
                      masking_key='ciao')
        pframe = p.execute(frame.msg)
        self.assertTrue(pframe)
        self.assertEqual(pframe.payload_length, len(self.large_bdata))
        self.assertEqual(pframe.body, self.large_bdata)
        
    def testPartialParsing(self):
        p = FrameParser()
        frame = Frame(self.large_bdata, opcode=0x2, final=True,
                      masking_key='ciao')
        self.assertEqual(p.execute(frame.msg[:1]), None)
        self.assertEqual(p.execute(frame.msg[1:5]), None)
        self.assertEqual(p.execute(frame.msg[5:50]), None)
        pframe = p.execute(frame.msg[50:])
        self.assertTrue(pframe)
        self.assertEqual(pframe.payload_length, len(self.large_bdata))
        self.assertEqual(pframe.body, self.large_bdata)