THRIFT-3516 Add feature test for THeader TBinaryProtocol

This closes #767

4 files changed, 112 insertions, 0 deletions

diff --git a/test/features/local_thrift/__init__.py b/test/features/local_thrift/__init__.py
new file mode 100644
index 000000000..383ee5f40
--- /dev/null
+++ b/test/features/local_thrift/__init__.py
@@ -0,0 +1,14 @@
+import os
+import sys
+
+SCRIPT_DIR = os.path.realpath(os.path.dirname(__file__))
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(SCRIPT_DIR)))
+
+if sys.version_info[0] == 2:
+  import glob
+  libdir = glob.glob(os.path.join(ROOT_DIR, 'lib', 'py', 'build', 'lib.*'))[0]
+  sys.path.insert(0, libdir)
+  thrift = __import__('thrift')
+else:
+  sys.path.insert(0, os.path.join(ROOT_DIR, 'lib', 'py', 'build', 'lib'))
+  thrift = __import__('thrift')
diff --git a/test/features/tests.json b/test/features/tests.json
new file mode 100644
index 000000000..0836309cb
--- /dev/null
+++ b/test/features/tests.json
@@ -0,0 +1,27 @@
+[
+  {
+    "description": "THeader detects unframed binary wire format",
+    "name": "theader_unframed_binary",
+    "command": [
+      "python",
+      "theader_binary.py"
+    ],
+    "protocols": ["header"],
+    "transports": ["buffered"],
+    "sockets": ["ip"],
+    "workdir": "features"
+  },
+  {
+    "description": "THeader detects framed binary wire format",
+    "name": "theader_framed_binary",
+    "command": [
+      "python",
+      "theader_binary.py",
+      "--override-transport=framed"
+    ],
+    "protocols": ["header"],
+    "transports": ["buffered"],
+    "sockets": ["ip"],
+    "workdir": "features"
+  }
+]
diff --git a/test/features/theader_binary.py b/test/features/theader_binary.py
new file mode 100644
index 000000000..0316741bc
--- /dev/null
+++ b/test/features/theader_binary.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+import argparse
+import socket
+import sys
+
+from util import add_common_args
+from local_thrift import thrift
+from thrift.Thrift import TMessageType, TType
+from thrift.transport.TSocket import TSocket
+from thrift.transport.TTransport import TBufferedTransport, TFramedTransport
+from thrift.protocol.TBinaryProtocol import TBinaryProtocol
+
+
+# THeader stack should accept binary protocol with optionally framed transport
+def main(argv):
+  p = argparse.ArgumentParser()
+  add_common_args(p)
+  # Since THeaderTransport acts as framed transport when detected frame, we
+  # cannot use --transport=framed as it would result in 2 layered frames.
+  p.add_argument('--override-transport')
+  args = p.parse_args()
+  assert args.protocol == 'header'
+  assert args.transport == 'buffered'
+  assert not args.ssl
+
+  sock = TSocket(args.host, args.port, socket_family=socket.AF_INET)
+  if not args.override_transport or args.override_transport == 'buffered':
+    trans = TBufferedTransport(sock)
+  elif args.override_transport == 'framed':
+    trans = TFramedTransport(sock)
+  else:
+    raise ValueError('invalid transport')
+  trans.open()
+  proto = TBinaryProtocol(trans)
+  proto.writeMessageBegin('testVoid', TMessageType.CALL, 3)
+  proto.writeStructBegin('testVoid_args')
+  proto.writeFieldStop()
+  proto.writeStructEnd()
+  proto.writeMessageEnd()
+  trans.flush()
+
+  _, mtype, _ = proto.readMessageBegin()
+  assert mtype == TMessageType.REPLY
+  proto.readStructBegin()
+  _, ftype, _ = proto.readFieldBegin()
+  assert ftype == TType.STOP
+  proto.readFieldEnd()
+  proto.readStructEnd()
+  proto.readMessageEnd()
+
+  trans.close()
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv[1:]))
diff --git a/test/features/util.py b/test/features/util.py
new file mode 100644
index 000000000..cff7ff8a2
--- /dev/null
+++ b/test/features/util.py
@@ -0,0 +1,15 @@
+import argparse
+
+
+def add_common_args(p):
+  p.add_argument('--host', default='localhost')
+  p.add_argument('--port', type=int)
+  p.add_argument('--protocol')
+  p.add_argument('--transport')
+  p.add_argument('--ssl', action='store_true')
+
+
+def parse_common_args(argv):
+  p = argparse.ArgumentParser()
+  add_common_args(p)
+  return p.parse_args(argv)