text: # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  # SPDX-License-Identifier: Apache-2.0.
  
  from awscrt.io import ClientBootstrap, ClientTlsContext, DefaultHostResolver, EventLoopGroup, Pkcs11Lib, TlsContextOptions, LogLevel, init_logging
  from awscrt.mqtt import Client, Connection, QoS
  from test import NativeResourceTest
  from concurrent.futures import Future
  import enum
  import os
  import pathlib
  import sys
  import unittest
  import uuid
  
  TIMEOUT = 100.0
  
  
  class MqttClientTest(NativeResourceTest):
      def test_lifetime(self):
          elg = EventLoopGroup()
          resolver = DefaultHostResolver(elg)
          bootstrap = ClientBootstrap(elg, resolver)
          client = Client(bootstrap)
  
  
  AuthType = enum.Enum('AuthType', ['CERT_AND_KEY', 'PKCS11'])
  
  
  class Config:
      def __init__(self, auth_type):
          self.endpoint = self._get_env('AWS_TEST_IOT_MQTT_ENDPOINT')
          self.cert_path = self._get_env('AWS_TEST_TLS_CERT_PATH')
          self.cert = pathlib.Path(self.cert_path).read_text().encode('utf-8')
  
          if auth_type == AuthType.CERT_AND_KEY:
              self.key_path = self._get_env('AWS_TEST_TLS_KEY_PATH')
              self.key = pathlib.Path(self.key_path).read_text().encode('utf-8')
  
          elif auth_type == AuthType.PKCS11:
              self.pkcs11_lib_path = self._get_env('AWS_TEST_PKCS11_LIB')
              self.pkcs11_pin = self._get_env('AWS_TEST_PKCS11_PIN')
              self.pkcs11_token_label = self._get_env('AWS_TEST_PKCS11_TOKEN_LABEL')
              self.pkcs11_key_label = self._get_env('AWS_TEST_PKCS11_KEY_LABEL')
  
      def _get_env(self, name):
          val = os.environ.get(name)
          if not val:
              raise unittest.SkipTest(f"test requires env var: {name}")
          return val
  
  
  def create_client_id():
      return 'aws-crt-python-unit-test-{0}'.format(uuid.uuid4())
  
  
  class MqttConnectionTest(NativeResourceTest):
      TEST_TOPIC = '/test/me/senpai'
      TEST_MSG = 'NOTICE ME!'.encode('utf8')
  
      def _create_connection(self, auth_type=AuthType.CERT_AND_KEY, use_static_singletons=False):
          config = Config(auth_type)
  
          if auth_type == AuthType.CERT_AND_KEY:
              tls_opts = TlsContextOptions.create_client_with_mtls_from_path(config.cert_path, config.key_path)
              tls = ClientTlsContext(tls_opts)
  
          elif auth_type == AuthType.PKCS11:
              try:
                  pkcs11_lib = Pkcs11Lib(
                      file=config.pkcs11_lib_path,
                      behavior=Pkcs11Lib.InitializeFinalizeBehavior.STRICT)
  
                  tls_opts = TlsContextOptions.create_client_with_mtls_pkcs11(
                      pkcs11_lib=pkcs11_lib,
                      user_pin=config.pkcs11_pin,
                      token_label=config.pkcs11_token_label,
                      private_key_label=config.pkcs11_key_label,
                      cert_file_path=config.cert_path)
  
                  tls = ClientTlsContext(tls_opts)
  
              except Exception as e:
                  if 'AWS_ERROR_UNIMPLEMENTED' in str(e):
                      raise unittest.SkipTest(f'TLS with PKCS#11 not supported on this platform ({sys.platform})')
                  else:
                      # re-raise exception
                      raise
  
          if use_static_singletons:
              client = Client(tls_ctx=tls)
          else:
              elg = EventLoopGroup()
              resolver = DefaultHostResolver(elg)
              bootstrap = ClientBootstrap(elg, resolver)
              client = Client(bootstrap, tls)
  
          connection = Connection(
              client=client,
              client_id=create_client_id(),
              host_name=config.endpoint,
              port=8883)
          return connection
  
      def test_connect_disconnect(self):
          connection = self._create_connection()
          connection.connect().result(TIMEOUT)
          connection.disconnect().result(TIMEOUT)
  
      def test_pkcs11(self):
          connection = self._create_connection(AuthType.PKCS11)
          connection.connect().result(TIMEOUT)
          connection.disconnect().result(TIMEOUT)
  
      def test_pub_sub(self):
          connection = self._create_connection()
          connection.connect().result(TIMEOUT)
          received = Future()
  
          def on_message(**kwargs):
              received.set_result(kwargs)
  
          # subscribe
          subscribed, packet_id = connection.subscribe(self.TEST_TOPIC, QoS.AT_LEAST_ONCE, on_message)
          suback = subscribed.result(TIMEOUT)
          self.assertEqual(packet_id, suback['packet_id'])
          self.assertEqual(self.TEST_TOPIC, suback['topic'])
          self.assertIs(QoS.AT_LEAST_ONCE, suback['qos'])
  
          # publish
          published, packet_id = connection.publish(self.TEST_TOPIC, self.TEST_MSG, QoS.AT_LEAST_ONCE)
          puback = published.result(TIMEOUT)
          self.assertEqual(packet_id, puback['packet_id'])
  
          # receive message
          rcv = received.result(TIMEOUT)
          self.assertEqual(self.TEST_TOPIC, rcv['topic'])
          self.assertEqual(self.TEST_MSG, rcv['payload'])
          self.assertFalse(rcv['dup'])
          self.assertEqual(QoS.AT_LEAST_ONCE, rcv['qos'])
          self.assertFalse(rcv['retain'])
  
          # unsubscribe
          unsubscribed, packet_id = connection.unsubscribe(self.TEST_TOPIC)
          unsuback = unsubscribed.result(TIMEOUT)
          self.assertEqual(packet_id, unsuback['packet_id'])
  
          # disconnect
          connection.disconnect().result(TIMEOUT)
  
      def test_on_message(self):
          connection = self._create_connection()
          received = Future()
  
          def on_message(**kwargs):
              received.set_result(kwargs)
  
          # on_message for connection has to be set before connect, or possible race will happen
          connection.on_message(on_message)
  
          connection.connect().result(TIMEOUT)
          # subscribe without callback
          subscribed, packet_id = connection.subscribe(self.TEST_TOPIC, QoS.AT_LEAST_ONCE)
          subscribed.result(TIMEOUT)
  
          # publish
          published, packet_id = connection.publish(self.TEST_TOPIC, self.TEST_MSG, QoS.AT_LEAST_ONCE)
          puback = published.result(TIMEOUT)
  
          # receive message
          rcv = received.result(TIMEOUT)
          self.assertEqual(self.TEST_TOPIC, rcv['topic'])
          self.assertEqual(self.TEST_MSG, rcv['payload'])
          self.assertFalse(rcv['dup'])
          self.assertEqual(QoS.AT_LEAST_ONCE, rcv['qos'])
          self.assertFalse(rcv['retain'])
  
          # disconnect
          connection.disconnect().result(TIMEOUT)
  
      def test_on_message_old_fn_signature(self):
          # ensure that message-received callbacks with the old function signature still work
          connection = self._create_connection()
  
          any_received = Future()
          sub_received = Future()
  
          # Note: Testing degenerate callback signature that failed to take
          # forward-compatibility **kwargs.
          def on_any_message(topic, payload):
              any_received.set_result({'topic': topic, 'payload': payload})
  
          def on_sub_message(topic, payload):
              sub_received.set_result({'topic': topic, 'payload': payload})
  
          # on_message for connection has to be set before connect, or possible race will happen
          connection.on_message(on_any_message)
  
          connection.connect().result(TIMEOUT)
          # subscribe without callback
          subscribed, packet_id = connection.subscribe(self.TEST_TOPIC, QoS.AT_LEAST_ONCE, on_sub_message)
          subscribed.result(TIMEOUT)
  
          # publish
          published, packet_id = connection.publish(self.TEST_TOPIC, self.TEST_MSG, QoS.AT_LEAST_ONCE)
          puback = published.result(TIMEOUT)
  
          # receive message
          rcv = any_received.result(TIMEOUT)
          self.assertEqual(self.TEST_TOPIC, rcv['topic'])
          self.assertEqual(self.TEST_MSG, rcv['payload'])
  
          rcv = sub_received.result(TIMEOUT)
          self.assertEqual(self.TEST_TOPIC, rcv['topic'])
          self.assertEqual(self.TEST_MSG, rcv['payload'])
  
          # disconnect
          connection.disconnect().result(TIMEOUT)
  
      def test_connect_disconnect_with_default_singletons(self):
          connection = self._create_connection(use_static_singletons=True)
          connection.connect().result(TIMEOUT)
          connection.disconnect().result(TIMEOUT)
  
          # free singletons
          ClientBootstrap.release_static_default()
          EventLoopGroup.release_static_default()
          DefaultHostResolver.release_static_default()
  
  
  if __name__ == 'main':
      unittest.main()
  