#!/usr/bin/python3
import socket
import unittest

import os
HOST = os.getenv('WEBDIS_HOST', '127.0.0.1')
PORT = int(os.getenv('WEBDIS_PORT', 7379))

class BlockingSocket:

	def __init__(self):
		self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		self.s.setblocking(True)
		self.s.connect((HOST, PORT))
	
	def __del__(self):
		self.s.close()

	def recv(self):
		out = b""
		while True:
			try:
				ret = self.s.recv(4096)
			except:
				return out
			if len(ret) == 0:
				return out
			out += ret

	def recv_until(self, limit):
		out = b""
		while not limit in out:
			try:
				out += self.s.recv(4096)
			except:
				return False
		return out

	def send(self, buf):
		sz = len(buf)
		done = 0
		while done < sz:
			try:
				ret = self.s.send(buf[done:4096])
			except:
				return False
			done = done + ret
			# print "Sent %d/%d so far (%s just now)" % (done, sz, ret)
			if ret < 0:
				return False
		return True

class LargeString:

	def __init__(self, c, n):
		self.char = c
		self.len = n

	def __len__(self):
		return self.len

	def __getitem__(self, chunk):
		if chunk.start > self.len:
			return b""
		if chunk.start + chunk.stop > self.len:
			return self.char * (self.len - chunk.start)
		return self.char * chunk.stop

class TestSocket(unittest.TestCase):

	def __init__(self, *args, **kwargs):
		super(TestSocket, self).__init__(*args, **kwargs)
		self.s = BlockingSocket()


class TestHugeUrl(TestSocket):

	def test_huge_url(self):
		n = 1024*1024*1024	# 1GB query-string

		start = b"GET /GET/x"
		end = b" HTTP/1.0\r\n\r\n"

		ok = self.s.send(start)
		fail1 = self.s.send(LargeString(b"A", n))
		fail2 = self.s.send(end)
		out = self.s.recv()

		self.assertTrue(ok)
		self.assertTrue(b"400 Bad Request" in out)

	def test_huge_upload(self):
		n = 1024*1024*1024	# upload 1GB

		start = b"PUT /SET/x HTTP/1.0\r\n"\
		+ ("Content-Length: %d\r\n" % (n)).encode('utf-8')\
		+ b"Expect: 100-continue\r\n\r\n"

		ok = self.s.send(start)
		cont = self.s.recv_until(b"\r\n")
		fail = self.s.send(LargeString(b"A", n))

		self.assertTrue(ok)
		self.assertTrue(b"HTTP/1.1 100 Continue" in cont)
		self.assertFalse(fail)

if __name__ == '__main__':
	unittest.main(verbosity=5)