1 files changed, 16 insertions, 5 deletions
diff --git a/conformance/conformance_python.py b/conformance/conformance_python.py
index 7ace9b16..c5ba2467 100755
--- a/conformance/conformance_python.py
+++ b/conformance/conformance_python.py
@@ -38,9 +38,12 @@ See conformance.proto for more information.
 import struct
 import sys
 import os
+from google.protobuf import descriptor
+from google.protobuf import descriptor_pool
 from google.protobuf import json_format
 from google.protobuf import message
 from google.protobuf import test_messages_proto3_pb2
+from google.protobuf import test_messages_proto2_pb2
 import conformance_pb2
 
 sys.stdout = os.fdopen(sys.stdout.fileno(), 'wb', 0)
@@ -53,9 +56,17 @@ class ProtocolError(Exception):
   pass
 
 def do_test(request):
-  test_message = test_messages_proto3_pb2.TestAllTypes()
+  isProto3 = (request.message_type == "protobuf_test_messages.proto3.TestAllTypesProto3")
+  isJson = (request.WhichOneof('payload') == 'json_payload')
+  isProto2 = (request.message_type == "protobuf_test_messages.proto2.TestAllTypesProto2")
+  
+  if (not isProto3) and (not isJson) and (not isProto2):
+    raise ProtocolError("Protobuf request doesn't have specific payload type")
+      
+  test_message = test_messages_proto2_pb2.TestAllTypesProto2() if isProto2 else \
+    test_messages_proto3_pb2.TestAllTypesProto3()
+
   response = conformance_pb2.ConformanceResponse()
-  test_message = test_messages_proto3_pb2.TestAllTypes()
 
   try:
     if request.WhichOneof('payload') == 'protobuf_payload':
@@ -63,8 +74,8 @@ def do_test(request):
         test_message.ParseFromString(request.protobuf_payload)
       except message.DecodeError as e:
         response.parse_error = str(e)
-        return response
-
+        return response  
+      
     elif request.WhichOneof('payload') == 'json_payload':
       try:
         json_format.Parse(request.json_payload, test_message)
@@ -82,7 +93,7 @@ def do_test(request):
       response.protobuf_payload = test_message.SerializeToString()
 
     elif request.requested_output_format == conformance_pb2.JSON:
-      try:
+      try: 
         response.json_payload = json_format.MessageToJson(test_message)
       except Exception as e:
         response.serialize_error = str(e)