python3-postgresql 1.1.0-1build1 / usr / lib / python3 / dist-packages / postgresql / test / test_protocol.py

##
# .test.test_protocol
##
import sys
import unittest
import struct
import decimal
import socket
import time
from threading import Thread

from ..protocol import element3 as e3
from ..protocol import xact3 as x3
from ..protocol import client3 as c3
from ..protocol import buffer as pq_buf
from ..python.socket import find_available_port, SocketFactory

def pair(msg):
	return (msg.type, msg.serialize())
def pairs(*msgseq):
	return list(map(pair, msgseq))

long = struct.Struct("!L")
packl = long.pack
unpackl = long.unpack

class test_buffer(unittest.TestCase):
	def setUp(self):
		self.buffer = pq_buf.pq_message_stream()

	def testMultiByteMessage(self):
		b = self.buffer
		b.write(b's')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x00\x00')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x00\x10')
		self.assertTrue(b.next_message() is None)
		data = b'twelve_chars'
		b.write(data)
		self.assertEqual(b.next_message(), (b's', data))

	def testSingleByteMessage(self):
		b = self.buffer
		b.write(b's')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x00')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x00\x00\x05')
		self.assertTrue(b.next_message() is None)
		b.write(b'b')
		self.assertEqual(b.next_message(), (b's', b'b'))

	def testEmptyMessage(self):
		b = self.buffer
		b.write(b'x')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x00\x00\x00')
		self.assertTrue(b.next_message() is None)
		b.write(b'\x04')
		self.assertEqual(b.next_message(), (b'x', b''))

	def testInvalidLength(self):
		b = self.buffer
		b.write(b'y\x00\x00\x00\x03')
		self.assertRaises(ValueError, b.next_message,)

	def testRemainder(self):
		b = self.buffer
		b.write(b'r\x00\x00\x00\x05Aremainder')
		self.assertEqual(b.next_message(), (b'r', b'A'))

	def testLarge(self):
		b = self.buffer
		factor = 1024
		r = 10000
		b.write(b'X' + packl(factor * r + 4))
		segment = b'\x00' * factor
		for x in range(r-1):
			b.write(segment)
		b.write(segment)
		msg = b.next_message()
		self.assertTrue(msg is not None)
		self.assertEqual(msg[0], b'X')

	def test_getvalue(self):
		# Make sure that getvalue() only applies to messages
		# that have not been read.
		b = self.buffer
		# It should be empty.
		self.assertEqual(b.getvalue(), b'')
		d = b'F' + packl(28)
		b.write(d)
		self.assertEqual(b.getvalue(), d)
		d1 = b'01'*12 # 24
		b.write(d1)
		self.assertEqual(b.getvalue(), d + d1)
		out = b.read()[0]
		self.assertEqual(out, (b'F', d1))
		nd = b'N'
		b.write(nd)
		self.assertEqual(b.getvalue(), nd)
		b.write(packl(4))
		self.assertEqual(list(b.read()), [(b'N', b'')])
		self.assertEqual(b.getvalue(), b'')
		# partial; read one message to exercise
		# that the appropriate fragment of the first
		# chunk in the buffer is picked up.
		first_body = (b'1234' * 3)
		first = b'v' + packl(len(first_body) + 4) + first_body
		second_body = (b'4321' * 5)
		second = b'z' + packl(len(second_body) + 4) + second_body
		b.write(first + second)
		self.assertEqual(b.getvalue(), first + second)
		self.assertEqual(list(b.read(1)), [(b'v', first_body)])
		self.assertEqual(b.getvalue(), second)
		self.assertEqual(list(b.read(1)), [(b'z', second_body)])
		# now, with a third full message in the next chunk
		third_body = (b'9876' * 10)
		third = b'3' + packl(len(third_body) + 4) + third_body
		b.write(first + second)
		b.write(third)
		self.assertEqual(b.getvalue(), first + second + third)
		self.assertEqual(list(b.read(1)), [(b'v', first_body)])
		self.assertEqual(b.getvalue(), second + third)
		self.assertEqual(list(b.read(1)), [(b'z', second_body)])
		self.assertEqual(b.getvalue(), third)
		self.assertEqual(list(b.read(1)), [(b'3', third_body)])
		self.assertEqual(b.getvalue(), b'')

##
# element3 tests
##

message_samples = [
	e3.VoidMessage,
	e3.Startup([
		(b'user', b'jwp'),
		(b'database', b'template1'),
		(b'options', b'-f'),
	]),
	e3.Notice((
		(b'S', b'FATAL'),
		(b'M', b'a descriptive message'),
		(b'C', b'FIVEC'),
		(b'D', b'bleh'),
		(b'H', b'dont spit into the fan'),
	)),
	e3.Notify(123, b'wood_table'),
	e3.KillInformation(19320, 589483),
	e3.ShowOption(b'foo', b'bar'),
	e3.Authentication(4, b'salt'),
	e3.Complete(b'SELECT'),
	e3.Ready(b'I'),
	e3.CancelRequest(4123, 14252),
	e3.NegotiateSSL(),
	e3.Password(b'ckr4t'),
	e3.AttributeTypes(()),
	e3.AttributeTypes(
		(123,) * 1
	),
	e3.AttributeTypes(
		(123,0) * 1
	),
	e3.AttributeTypes(
		(123,0) * 2
	),
	e3.AttributeTypes(
		(123,0) * 4
	),
	e3.TupleDescriptor(()),
	e3.TupleDescriptor((
		(b'name', 123, 1, 1, 0, 0, 1,),
	)),
	e3.TupleDescriptor((
		(b'name', 123, 1, 2, 0, 0, 1,),
	) * 2),
	e3.TupleDescriptor((
		(b'name', 123, 1, 2, 1, 0, 1,),
	) * 3),
	e3.TupleDescriptor((
		(b'name', 123, 1, 1, 0, 0, 1,),
	) * 1000),
	e3.Tuple([]),
	e3.Tuple([b'foo',]),
	e3.Tuple([None]),
	e3.Tuple([b'foo',b'bar']),
	e3.Tuple([None, None]),
	e3.Tuple([None, b'foo', None]),
	e3.Tuple([b'bar', None, b'foo', None, b'bleh']),
	e3.Tuple([b'foo', b'bar'] * 100),
	e3.Tuple([None] * 100),
	e3.Query(b'select * from u'),
	e3.Parse(b'statement_id', b'query', (123, 0)),
	e3.Parse(b'statement_id', b'query', (123,)),
	e3.Parse(b'statement_id', b'query', ()),
	e3.Bind(b'portal_id', b'statement_id',
		(b'tt',b'\x00\x00'),
		[b'data',None], (b'ff',b'xx')),
	e3.Bind(b'portal_id', b'statement_id', (b'tt',), [None], (b'xx',)),
	e3.Bind(b'portal_id', b'statement_id', (b'ff',), [b'data'], ()),
	e3.Bind(b'portal_id', b'statement_id', (), [], (b'xx',)),
	e3.Bind(b'portal_id', b'statement_id', (), [], ()),
	e3.Execute(b'portal_id', 500),
	e3.Execute(b'portal_id', 0),
	e3.DescribeStatement(b'statement_id'),
	e3.DescribePortal(b'portal_id'),
	e3.CloseStatement(b'statement_id'),
	e3.ClosePortal(b'portal_id'),
	e3.Function(123, (), [], b'xx'),
	e3.Function(321, (b'tt',), [b'foo'], b'xx'),
	e3.Function(321, (b'tt',), [None], b'xx'),
	e3.Function(321, (b'aa', b'aa'), [None,b'a' * 200], b'xx'),
	e3.FunctionResult(b''),
	e3.FunctionResult(b'foobar'),
	e3.FunctionResult(None),
	e3.CopyToBegin(123, [321,123]),
	e3.CopyToBegin(0, [10,]),
	e3.CopyToBegin(123, []),
	e3.CopyFromBegin(123, [321,123]),
	e3.CopyFromBegin(0, [10]),
	e3.CopyFromBegin(123, []),
	e3.CopyData(b''),
	e3.CopyData(b'foo'),
	e3.CopyData(b'a' * 2048),
	e3.CopyFail(b''),
	e3.CopyFail(b'iiieeeeee!'),
]

class test_element3(unittest.TestCase):
	def test_cat_messages(self):
		# The optimized implementation will identify adjacent copy data, and
		# take a more efficient route; so rigorously test the switch between the
		# two modes.
		self.assertEqual(e3.cat_messages([]), b'')
		self.assertEqual(e3.cat_messages([b'foo']), b'd\x00\x00\x00\x07foo')
		self.assertEqual(e3.cat_messages([b'foo', b'foo']), 2*b'd\x00\x00\x00\x07foo')
		# copy, other, copy
		self.assertEqual(e3.cat_messages([b'foo', e3.SynchronizeMessage, b'foo']),
			b'd\x00\x00\x00\x07foo' + e3.SynchronizeMessage.bytes() + b'd\x00\x00\x00\x07foo')
		# copy, other, copy*1000
		self.assertEqual(e3.cat_messages(1000*[b'foo', e3.SynchronizeMessage, b'foo']),
			1000*(b'd\x00\x00\x00\x07foo' + e3.SynchronizeMessage.bytes() + b'd\x00\x00\x00\x07foo'))
		# other, copy, copy*1000
		self.assertEqual(e3.cat_messages(1000*[e3.SynchronizeMessage, b'foo', b'foo']),
			1000*(e3.SynchronizeMessage.bytes() + 2*b'd\x00\x00\x00\x07foo'))
		pack_head = struct.Struct("!lH").pack
		# tuple
		self.assertEqual(e3.cat_messages([(b'foo',),]),
			b'D' + pack_head(7 + 4 + 2, 1) + b'\x00\x00\x00\x03foo')
		# tuple(foo,\N)
		self.assertEqual(e3.cat_messages([(b'foo',None,),]),
			b'D' + pack_head(7 + 4 + 4 + 2, 2) + b'\x00\x00\x00\x03foo\xFF\xFF\xFF\xFF')
		# tuple(foo,\N,bar)
		self.assertEqual(e3.cat_messages([(b'foo',None,b'bar'),]),
			b'D' + pack_head(7 + 7 + 4 + 4 + 2, 3) + \
			b'\x00\x00\x00\x03foo\xFF\xFF\xFF\xFF\x00\x00\x00\x03bar')
		# too many attributes
		self.assertRaises((OverflowError, struct.error),
			e3.cat_messages, [(None,) * 0x10000])

		class ThisEx(Exception):
			pass
		class ThatEx(Exception):
			pass
		class Bad(e3.Message):
			def serialize(self):
				raise ThisEx('foo')
		self.assertRaises(ThisEx, e3.cat_messages, [Bad()])
		class NoType(e3.Message):
			def serialize(self):
				return b''
		self.assertRaises(AttributeError, e3.cat_messages, [NoType()])
		class BadType(e3.Message):
			type = 123
			def serialize(self):
				return b''
		self.assertRaises((TypeError,struct.error), e3.cat_messages, [BadType()])


	def testSerializeParseConsistency(self):
		for msg in message_samples:
			smsg = msg.serialize()
			self.assertEqual(msg, msg.parse(smsg))

	def testEmptyMessages(self):
		for x in e3.__dict__.values():
			if isinstance(x, e3.EmptyMessage):
				xtype = type(x)
				self.assertTrue(x is xtype())

	def testUnknownNoticeFields(self):
		N = e3.Notice.parse(b'\x00\x00Z\x00Xklsvdnvldsvkndvlsn\x00Pfoobar\x00Mmessage\x00')
		E = e3.Error.parse(b'Z\x00Xklsvdnvldsvkndvlsn\x00Pfoobar\x00Mmessage\x00\x00')
		self.assertEqual(N[b'M'], b'message')
		self.assertEqual(E[b'M'], b'message')
		self.assertEqual(N[b'P'], b'foobar')
		self.assertEqual(E[b'P'], b'foobar')
		self.assertEqual(len(N), 4)
		self.assertEqual(len(E), 4)

	def testCompleteExtracts(self):
		x = e3.Complete(b'FOO BAR 1321')
		self.assertEqual(x.extract_command(), b'FOO BAR')
		self.assertEqual(x.extract_count(), 1321)
		x = e3.Complete(b' CREATE  	TABLE 13210  ')
		self.assertEqual(x.extract_command(), b'CREATE  	TABLE')
		self.assertEqual(x.extract_count(), 13210)
		x = e3.Complete(b'  CREATE  	TABLE  \t713210  ')
		self.assertEqual(x.extract_command(), b'CREATE  	TABLE')
		self.assertEqual(x.extract_count(), 713210)
		x = e3.Complete(b'  CREATE  	TABLE  0 \t13210  ')
		self.assertEqual(x.extract_command(), b'CREATE  	TABLE')
		self.assertEqual(x.extract_count(), 13210)
		x = e3.Complete(b' 0 \t13210 ')
		self.assertEqual(x.extract_command(), None)
		self.assertEqual(x.extract_count(), 13210)

##
# .protocol.xact3 tests
##

xact_samples = [
	# Simple contrived exchange.
	(
		(
			e3.Query(b"COMPLETE"),
		), (
			e3.Complete(b'COMPLETE'),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Query(b"ROW DATA"),
		), (
			e3.TupleDescriptor((
				(b'foo', 1, 1, 1, 1, 1, 1),
				(b'bar', 1, 2, 1, 1, 1, 1),
			)),
			e3.Tuple((b'lame', b'lame')),
			e3.Complete(b'COMPLETE'),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Query(b"ROW DATA"),
		), (
			e3.TupleDescriptor((
				(b'foo', 1, 1, 1, 1, 1, 1),
				(b'bar', 1, 2, 1, 1, 1, 1),
			)),
			e3.Tuple((b'lame', b'lame')),
			e3.Tuple((b'lame', b'lame')),
			e3.Tuple((b'lame', b'lame')),
			e3.Tuple((b'lame', b'lame')),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Query(b"NULL"),
		), (
			e3.Null(),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Query(b"COPY TO"),
		), (
			e3.CopyToBegin(1, [1,2]),
			e3.CopyData(b'row1'),
			e3.CopyData(b'row2'),
			e3.CopyDone(),
			e3.Complete(b'COPY TO'),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Function(1, [b''], [b''], 1),
		), (
			e3.FunctionResult(b'foo'),
			e3.Ready(b'I'),
		)
	),
	(
		(
			e3.Parse(b"NAME", b"SQL", ()),
		), (
			e3.ParseComplete(),
		)
	),
	(
		(
			e3.Bind(b"NAME", b"STATEMENT_ID", (), (), ()),
		), (
			e3.BindComplete(),
		)
	),
	(
		(
			e3.Parse(b"NAME", b"SQL", ()),
			e3.Bind(b"NAME", b"STATEMENT_ID", (), (), ()),
		), (
			e3.ParseComplete(),
			e3.BindComplete(),
		)
	),
	(
		(
			e3.Describe(b"STATEMENT_ID"),
		), (
			e3.AttributeTypes(()),
			e3.NoData(),
		)
	),
	(
		(
			e3.Describe(b"STATEMENT_ID"),
		), (
			e3.AttributeTypes(()),
			e3.TupleDescriptor(()),
		)
	),
	(
		(
			e3.CloseStatement(b"foo"),
		), (
			e3.CloseComplete(),
		),
	),
	(
		(
			e3.ClosePortal(b"foo"),
		), (
			e3.CloseComplete(),
		),
	),
	(
		(
			e3.Synchronize(),
		), (
			e3.Ready(b'I'),
		),
	),
]

class test_xact3(unittest.TestCase):
	def testTransactionSamplesAll(self):
		for xcmd, xres in xact_samples:
			x = x3.Instruction(xcmd)
			r = tuple([(y.type, y.serialize()) for y in xres])
			x.state[1]()
			self.assertEqual(x.messages, ())
			x.state[1](r)
			self.assertEqual(x.state, x3.Complete)
			rec = []
			for y in x.completed:
				for z in y[1]:
					if type(z) is type(b''):
						z = e3.CopyData(z)
					rec.append(z)
			self.assertEqual(xres, tuple(rec))

	def testClosing(self):
		c = x3.Closing()
		self.assertEqual(c.messages, (e3.DisconnectMessage,))
		c.state[1]()
		self.assertEqual(c.fatal, True)
		self.assertEqual(c.error_message.__class__, e3.ClientError)
		self.assertEqual(c.error_message[b'C'], '08003')

	def testNegotiation(self):
		# simple successful run
		n = x3.Negotiation({}, b'')
		n.state[1]()
		n.state[1](
			pairs(
				e3.Notice(((b'M', b"foobar"),)),
				e3.Authentication(e3.AuthRequest_OK, b''),
				e3.KillInformation(0,0),
				e3.ShowOption(b'name', b'val'),
				e3.Ready(b'I'),
			)
		)
		self.assertEqual(n.state, x3.Complete)
		self.assertEqual(n.last_ready.xact_state, b'I')
		# no killinfo.. should cause protocol error...
		n = x3.Negotiation({}, b'')
		n.state[1]()
		n.state[1](
			pairs(
				e3.Notice(((b'M', b"foobar"),)),
				e3.Authentication(e3.AuthRequest_OK, b''),
				e3.ShowOption(b'name', b'val'),
				e3.Ready(b'I'),
			)
		)
		self.assertEqual(n.state, x3.Complete)
		self.assertEqual(n.last_ready, None)
		self.assertEqual(n.error_message[b'C'], '08P01')
		# killinfo twice.. must cause protocol error...
		n = x3.Negotiation({}, b'')
		n.state[1]()
		n.state[1](
			pairs(
				e3.Notice(((b'M', b"foobar"),)),
				e3.Authentication(e3.AuthRequest_OK, b''),
				e3.ShowOption(b'name', b'val'),
				e3.KillInformation(0,0),
				e3.KillInformation(0,0),
				e3.Ready(b'I'),
			)
		)
		self.assertEqual(n.state, x3.Complete)
		self.assertEqual(n.last_ready, None)
		self.assertEqual(n.error_message[b'C'], '08P01')
		# start with ready message..
		n = x3.Negotiation({}, b'')
		n.state[1]()
		n.state[1](
			pairs(
				e3.Notice(((b'M', b"foobar"),)),
				e3.Ready(b'I'),
				e3.Authentication(e3.AuthRequest_OK, b''),
				e3.ShowOption(b'name', b'val'),
			)
		)
		self.assertEqual(n.state, x3.Complete)
		self.assertEqual(n.last_ready, None)
		self.assertEqual(n.error_message[b'C'], '08P01')
		# unsupported authreq
		n = x3.Negotiation({}, b'')
		n.state[1]()
		n.state[1](
			pairs(
				e3.Authentication(255, b''),
			)
		)
		self.assertEqual(n.state, x3.Complete)
		self.assertEqual(n.last_ready, None)
		self.assertEqual(n.error_message[b'C'], '--AUT')

	def testInstructionAsynchook(self):
		l = []
		def hook(data):
			l.append(data)
		x = x3.Instruction([
			e3.Query(b"NOTHING")
		], asynchook = hook)
		a1 = e3.Notice(((b'M', b"m1"),))
		a2 = e3.Notify(0, b'relation', b'parameter')
		a3 = e3.ShowOption(b'optname', b'optval')
		# "send" the query message
		x.state[1]()
		# "receive" the tuple
		x.state[1]([(a1.type, a1.serialize()),])
		a2l = [(a2.type, a2.serialize()),]
		x.state[1](a2l)
		# validate that the hook is not fed twice because
		# it's the exact same message set. (later assertion will validate)
		x.state[1](a2l)
		x.state[1]([(a3.type, a3.serialize()),])
		# we only care about validating that l got everything.
		self.assertEqual([a1,a2,a3], l)
		self.assertEqual(x.state[0], x3.Receiving)
		# validate that the asynchook exception is trapped.
		class Nee(Exception):
			pass
		def ehook(msg):
			raise Nee("this should **not** be part of the summary")
		x = x3.Instruction([
			e3.Query(b"NOTHING")
		], asynchook = ehook)
		a1 = e3.Notice(((b'M', b"m1"),))
		x.state[1]()
		import sys
		v = None
		def exchook(typ, val, tb):
			nonlocal v
			v = val
		seh = sys.excepthook
		sys.excepthook = exchook
		# we only care about validating that the exchook got called.
		x.state[1]([(a1.type, a1.serialize())])
		sys.excepthook = seh
		self.assertTrue(isinstance(v, Nee))

class test_client3(unittest.TestCase):
	def test_timeout(self):
		portnum = find_available_port()
		servsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		with servsock:
			servsock.bind(('localhost', portnum))
			pc = c3.Connection(
				SocketFactory(
					(socket.AF_INET, socket.SOCK_STREAM),
					('localhost', portnum)
				),
				{}
			)
			pc.connect(timeout = 1)
			try:
				self.assertEqual(pc.xact.fatal, True)
				self.assertEqual(pc.xact.__class__, x3.Negotiation)
			finally:
				if pc.socket is not None:
					pc.socket.close()

	def test_SSL_failure(self):
		portnum = find_available_port()
		servsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		with servsock:
			servsock.bind(('localhost', portnum))
			pc = c3.Connection(
				SocketFactory(
					(socket.AF_INET, socket.SOCK_STREAM),
					('localhost', portnum)
				),
				{}
			)
			exc = None
			servsock.listen(1)
			def client_thread():
				pc.connect(ssl = True)
			client = Thread(target = client_thread)
			try:
				client.start()
				c, addr = servsock.accept()
				with c:
					c.send(b'S')
					c.sendall(b'0000000000000000000000')
					c.recv(1024)
					c.close()
				client.join()
			finally:
				if pc.socket is not None:
					pc.socket.close()

		self.assertEqual(pc.xact.fatal, True)
		self.assertEqual(pc.xact.__class__, x3.Negotiation)
		self.assertEqual(pc.xact.error_message.__class__, e3.ClientError)
		self.assertTrue(hasattr(pc.xact, 'exception'))

	def test_bad_negotiation(self):
		portnum = find_available_port()
		servsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		servsock.bind(('localhost', portnum))
		pc = c3.Connection(
			SocketFactory(
				(socket.AF_INET, socket.SOCK_STREAM),
				('localhost', portnum)
			),
			{}
		)
		exc = None
		servsock.listen(1)
		def client_thread():
			pc.connect()
		client = Thread(target = client_thread)
		try:
			client.start()
			c, addr = servsock.accept()
			try:
				c.recv(1024)
			finally:
				c.close()
			time.sleep(0.25)
			client.join()
			servsock.close()
			self.assertEqual(pc.xact.fatal, True)
			self.assertEqual(pc.xact.__class__, x3.Negotiation)
			self.assertEqual(pc.xact.error_message.__class__, e3.ClientError)
			self.assertEqual(pc.xact.error_message[b'C'], '08006')
		finally:
			servsock.close()
			if pc.socket is not None:
				pc.socket.close()

if __name__ == '__main__':
	from types import ModuleType
	this = ModuleType("this")
	this.__dict__.update(globals())
	try:
		unittest.main(this)
	finally:
		import gc
		gc.collect()