#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from ThriftTest.ttypes import Bonk, VersioningTestV1, VersioningTestV2
from thrift.protocol import TJSONProtocol
from thrift.transport import TTransport

import json
import unittest


class SimpleJSONProtocolTest(unittest.TestCase):
  protocol_factory = TJSONProtocol.TSimpleJSONProtocolFactory()

  def _assertDictEqual(self, a, b, msg=None):
    if hasattr(self, 'assertDictEqual'):
      # assertDictEqual only in Python 2.7. Depends on your machine.
      self.assertDictEqual(a, b, msg)
      return

    # Substitute implementation not as good as unittest library's
    self.assertEquals(len(a), len(b), msg)
    for k, v in a.iteritems():
      self.assertTrue(k in b, msg)
      self.assertEquals(b.get(k), v, msg)

  def _serialize(self, obj):
    trans = TTransport.TMemoryBuffer()
    prot = self.protocol_factory.getProtocol(trans)
    obj.write(prot)
    return trans.getvalue()

  def _deserialize(self, objtype, data):
    prot = self.protocol_factory.getProtocol(TTransport.TMemoryBuffer(data))
    ret = objtype()
    ret.read(prot)
    return ret

  def testWriteOnly(self):
    self.assertRaises(NotImplementedError,
                      self._deserialize, VersioningTestV1, b'{}')

  def testSimpleMessage(self):
      v1obj = VersioningTestV1(
          begin_in_both=12345,
          old_string='aaa',
          end_in_both=54321)
      expected = dict(begin_in_both=v1obj.begin_in_both,
                      old_string=v1obj.old_string,
                      end_in_both=v1obj.end_in_both)
      actual = json.loads(self._serialize(v1obj).decode('ascii'))

      self._assertDictEqual(expected, actual)

  def testComplicated(self):
      v2obj = VersioningTestV2(
          begin_in_both=12345,
          newint=1,
          newbyte=2,
          newshort=3,
          newlong=4,
          newdouble=5.0,
          newstruct=Bonk(message="Hello!", type=123),
          newlist=[7, 8, 9],
          newset=set([42, 1, 8]),
          newmap={1: 2, 2: 3},
          newstring="Hola!",
          end_in_both=54321)
      expected = dict(begin_in_both=v2obj.begin_in_both,
                      newint=v2obj.newint,
                      newbyte=v2obj.newbyte,
                      newshort=v2obj.newshort,
                      newlong=v2obj.newlong,
                      newdouble=v2obj.newdouble,
                      newstruct=dict(message=v2obj.newstruct.message,
                                     type=v2obj.newstruct.type),
                      newlist=v2obj.newlist,
                      newset=list(v2obj.newset),
                      newmap=v2obj.newmap,
                      newstring=v2obj.newstring,
                      end_in_both=v2obj.end_in_both)

      # Need to load/dump because map keys get escaped.
      expected = json.loads(json.dumps(expected))
      actual = json.loads(self._serialize(v2obj).decode('ascii'))
      self._assertDictEqual(expected, actual)


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