Add MQTT connect and publish unit tests (#964)

* Add unit tests for MQTT_Connect, MQTT_Publish, MQTT_Disconnect, MQTT_GetPacketId

* Update MQTT_GetPacketId for NULL checks

* Update MQTT_Connect for parameter checks

src/mqtt.c

@@ -940,12 +940,13 @@ MQTTStatus_t MQTT_Connect( MQTTContext_t * const pContext,
     MQTTStatus_t status = MQTTSuccess;
     MQTTPacketInfo_t incomingPacket = { .type = ( ( uint8_t ) 0 ) };
 
-    if( ( pContext == NULL ) || ( pConnectInfo == NULL ) )
+    if( ( pContext == NULL ) || ( pConnectInfo == NULL ) || ( pSessionPresent == NULL ) )
     {
         LogError( ( "Argument cannot be NULL: pContext=%p, "
-                    "pConnectInfo=%p.",
+                    "pConnectInfo=%p, pSessionPresent=%p.",
                     pContext,
-                    pConnectInfo ) );
+                    pConnectInfo,
+                    pSessionPresent ) );
         status = MQTTBadParameter;
     }
 
@@ -1458,13 +1459,16 @@ MQTTStatus_t MQTT_ProcessLoop( MQTTContext_t * const pContext,
 
 uint16_t MQTT_GetPacketId( MQTTContext_t * const pContext )
 {
-    uint16_t packetId = pContext->nextPacketId;
+    uint16_t packetId = 0U;
 
-    pContext->nextPacketId++;
+    if( pContext != NULL )
-
-    if( pContext->nextPacketId == 0U )
     {
-        pContext->nextPacketId = 1;
+        packetId = pContext->nextPacketId++;
+
+        if( pContext->nextPacketId == 0U )
+        {
+            pContext->nextPacketId++;
+        }
     }
 
     return packetId;

utest/CMakeLists.txt

@@ -10,6 +10,7 @@ set(project_name "mqtt")
 # list the files to mock here
 list(APPEND mock_list
             "${MODULES_DIR}/standard/mqtt/include/mqtt_lightweight.h"
+            "${MODULES_DIR}/standard/mqtt/include/mqtt_state.h"
         )
 # list the directories your mocks need
 list(APPEND mock_include_list

utest/mqtt_utest.c

@@ -1,18 +1,41 @@
+#include <string.h>
+#include <stdbool.h>
+
 #include "unity.h"
 
 /* Include paths for public enums, structures, and macros. */
 #include "mqtt.h"
 
+#include "mock_mqtt_lightweight.h"
+#include "mock_mqtt_state.h"
+
 /**
  * @brief A valid starting packet ID per MQTT spec. Start from 1.
  */
 #define MQTT_NEXT_PACKET_ID_START    ( 1 )
 
+/**
+ * @brief Length of the MQTT network buffer.
+ */
+#define MQTT_TEST_BUFFER_LENGTH      ( 1024 )
+
+/**
+ * @brief Time at the beginning of each test. Note that this is not updated with
+ * a real clock. Instead, we simply increment this variable.
+ */
+static uint32_t globalEntryTime = 0;
+
+/**
+ * @brief A static buffer used by the MQTT library for storing packet data.
+ */
+static uint8_t mqttBuffer[ MQTT_TEST_BUFFER_LENGTH ] = { 0 };
+
 /* ============================   UNITY FIXTURES ============================ */
 
 /* Called before each test method. */
 void setUp()
 {
+    memset( ( void * ) mqttBuffer, 0x0, sizeof( mqttBuffer ) );
 }
 
 /* Called after each test method. */
@@ -31,6 +54,150 @@ int suiteTearDown( int numFailures )
     return numFailures;
 }
 
+/* ========================================================================== */
+
+/**
+ * @brief Mock successful transport send, and write data into buffer for
+ * verification.
+ */
+static int32_t mockSend( MQTTNetworkContext_t context,
+                         const void * pMessage,
+                         size_t bytesToSend )
+{
+    const uint8_t * buffer = ( const uint8_t * ) pMessage;
+    /* Treat network context as pointer to buffer for mocking. */
+    uint8_t * mockNetwork = ( *( uint8_t ** ) context );
+    size_t bytesSent = 0;
+
+    while( bytesSent++ < bytesToSend )
+    {
+        /* Write single byte and advance buffer. */
+        *mockNetwork++ = *buffer++;
+    }
+    /* Move stream by bytes sent. */
+    ( *( uint8_t ** ) context ) = mockNetwork;
+
+    return bytesToSend;
+}
+
+/**
+ * @brief Initialize pNetworkBuffer using static buffer.
+ *
+ * @param[in] pNetworkBuffer Network buffer provided for the context.
+ */
+static void setupNetworkBuffer( MQTTFixedBuffer_t * const pNetworkBuffer )
+{
+    pNetworkBuffer->pBuffer = mqttBuffer;
+    pNetworkBuffer->size = MQTT_TEST_BUFFER_LENGTH;
+}
+
+/**
+ * @brief Mocked MQTT event callback.
+ */
+static void eventCallback( MQTTContext_t * pContext,
+                           MQTTPacketInfo_t * pPacketInfo,
+                           uint16_t packetIdentifier,
+                           MQTTPublishInfo_t * pPublishInfo )
+{
+    ( void ) pContext;
+    ( void ) pPacketInfo;
+    ( void ) packetIdentifier;
+    ( void ) pPublishInfo;
+}
+
+/**
+ * @brief A mocked timer query function that increments on every call. This
+ * guarantees that only a single iteration runs in the ProcessLoop for ease
+ * of testing.
+ */
+static uint32_t getTime( void )
+{
+    return globalEntryTime++;
+}
+
+/**
+ * @brief Mocked successful transport send.
+ */
+static int32_t transportSendSuccess( MQTTNetworkContext_t pContext,
+                                     const void * pBuffer,
+                                     size_t bytesToWrite )
+{
+    ( void ) pContext;
+    ( void ) pBuffer;
+    return bytesToWrite;
+}
+
+/**
+ * @brief Mocked failed transport send.
+ */
+static int32_t transportSendFailure( MQTTNetworkContext_t pContext,
+                                     const void * pBuffer,
+                                     size_t bytesToWrite )
+{
+    ( void ) pContext;
+    ( void ) pBuffer;
+    ( void ) bytesToWrite;
+    return -1;
+}
+
+/**
+ * @brief Mocked successful transport read.
+ */
+static int32_t transportRecvSuccess( MQTTNetworkContext_t pContext,
+                                     void * pBuffer,
+                                     size_t bytesToRead )
+{
+    ( void ) pContext;
+    ( void ) pBuffer;
+    return bytesToRead;
+}
+
+/**
+ * @brief Mocked failed transport read.
+ */
+static int32_t transportRecvFailure( MQTTNetworkContext_t pContext,
+                                     void * pBuffer,
+                                     size_t bytesToRead )
+{
+    ( void ) pContext;
+    ( void ) pBuffer;
+    ( void ) bytesToRead;
+    return -1;
+}
+
+/**
+ * @brief Mocked failed transport read.
+ */
+static int32_t transportRecvOneByte( MQTTNetworkContext_t pContext,
+                                     void * pBuffer,
+                                     size_t bytesToRead )
+{
+    ( void ) pContext;
+    ( void ) pBuffer;
+    return 1;
+}
+
+/**
+ * @brief Initialize the transport interface with the mocked functions for
+ * send and receive.
+ */
+static void setupTransportInterface( MQTTTransportInterface_t * pTransport )
+{
+    pTransport->networkContext = 0;
+    pTransport->send = transportSendSuccess;
+    pTransport->recv = transportRecvSuccess;
+}
+
+/**
+ * @brief Initialize our event and time callback with the mocked functions
+ * defined for the purposes this test.
+ */
+static void setupCallbacks( MQTTApplicationCallbacks_t * pCallbacks )
+{
+    pCallbacks->appCallback = eventCallback;
+    pCallbacks->getTime = getTime;
+}
+
 /* ============================   Testing MQTT_Init ========================= */
 
 /**
@@ -78,3 +245,419 @@ void test_MQTT_Init_Invalid_Params( void )
     mqttStatus = MQTT_Init( &context, &transport, &callbacks, NULL );
     TEST_ASSERT_EQUAL( MQTTBadParameter, mqttStatus );
 }
+
+/* ========================================================================== */
+
+/**
+ * @brief Test MQTT_Connect, except for receiving the CONNACK.
+ */
+void test_MQTT_Connect_sendConnect( void )
+{
+    MQTTContext_t mqttContext;
+    MQTTConnectInfo_t connectInfo;
+    MQTTPublishInfo_t willInfo;
+    uint32_t timeout = 2;
+    bool sessionPresent;
+    MQTTStatus_t status;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    MQTTPacketInfo_t incomingPacket;
+    size_t remainingLength, packetSize;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    /* Check parameters */
+    status = MQTT_Connect( NULL, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+    status = MQTT_Connect( &mqttContext, NULL, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, NULL );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    /* Empty connect info fails. */
+    MQTT_GetConnectPacketSize_ExpectAnyArgsAndReturn( MQTTBadParameter );
+    memset( ( void * ) &connectInfo, 0x0, sizeof( connectInfo ) );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    connectInfo.pClientIdentifier = MQTT_CLIENT_IDENTIFIER;
+    connectInfo.clientIdentifierLength = sizeof( MQTT_CLIENT_IDENTIFIER ) - 1;
+
+    MQTT_GetConnectPacketSize_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_SerializeConnect_ExpectAnyArgsAndReturn( MQTTNoMemory );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTNoMemory, status );
+
+    MQTT_SerializeConnect_IgnoreAndReturn( MQTTSuccess );
+
+    /* Transport send failed when sending CONNECT. */
+    /* Choose 10 bytes variable header + 1 byte payload for the remaining
+     * length of the CONNECT. The packet size needs to be nonzero for this test
+     * as that is the amount of bytes used in the call to send the packet. */
+    packetSize = 13;
+    remainingLength = 11;
+    mqttContext.transportInterface.send = transportSendFailure;
+    MQTT_GetConnectPacketSize_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetConnectPacketSize_IgnoreArg_pPacketSize();
+    MQTT_GetConnectPacketSize_IgnoreArg_pRemainingLength();
+    MQTT_GetConnectPacketSize_ReturnThruPtr_pPacketSize( &packetSize );
+    MQTT_GetConnectPacketSize_ReturnThruPtr_pRemainingLength( &remainingLength );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTSendFailed, status );
+
+    /* Send the CONNECT successfully. This provides branch coverage for sendPacket. */
+    mqttContext.transportInterface.send = transportSendSuccess;
+    MQTT_GetConnectPacketSize_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetConnectPacketSize_ReturnThruPtr_pPacketSize( &packetSize );
+    MQTT_GetConnectPacketSize_ReturnThruPtr_pRemainingLength( &remainingLength );
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTRecvFailed );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTRecvFailed, status );
+}
+
+/**
+ * @brief Test CONNACK reception in MQTT_Connect.
+ */
+void test_MQTT_Connect_receiveConnack( void )
+{
+    MQTTContext_t mqttContext;
+    MQTTConnectInfo_t connectInfo;
+    uint32_t timeout = 0;
+    bool sessionPresent;
+    MQTTStatus_t status;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    MQTTPacketInfo_t incomingPacket;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+    transport.recv = transportRecvFailure;
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    /* Everything before receiving the CONNACK should succeed. */
+    MQTT_SerializeConnect_IgnoreAndReturn( MQTTSuccess );
+    MQTT_GetConnectPacketSize_IgnoreAndReturn( MQTTSuccess );
+
+    /* Nothing received from transport interface. Set timeout to 2 for branch coverage. */
+    timeout = 2;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTNoDataAvailable );
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTNoDataAvailable );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTNoDataAvailable, status );
+
+    /* Did not receive a CONNACK. */
+    incomingPacket.type = MQTT_PACKET_TYPE_PINGRESP;
+    incomingPacket.remainingLength = 0;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTBadResponse, status );
+
+    /* Transport receive failure when receiving rest of packet. */
+    incomingPacket.type = MQTT_PACKET_TYPE_CONNACK;
+    incomingPacket.remainingLength = 2;
+    timeout = 2;
+    mqttContext.transportInterface.recv = transportRecvFailure;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTRecvFailed, status );
+
+    /* Bad response when deserializing CONNACK. */
+    mqttContext.transportInterface.recv = transportRecvSuccess;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    MQTT_DeserializeAck_ExpectAnyArgsAndReturn( MQTTBadResponse );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTBadResponse, status );
+}
+
+/**
+ * @brief Test error cases for MQTT_Connect when a timeout occurs or the packet
+ * needs to be discarded in MQTT_Connect.
+ */
+void test_MQTT_Connect_partial_receive()
+{
+    MQTTContext_t mqttContext;
+    MQTTConnectInfo_t connectInfo;
+    uint32_t timeout = 0;
+    bool sessionPresent;
+    MQTTStatus_t status;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    MQTTPacketInfo_t incomingPacket;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+    transport.recv = transportRecvOneByte;
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    /* Everything before receiving the CONNACK should succeed. */
+    MQTT_SerializeConnect_IgnoreAndReturn( MQTTSuccess );
+    MQTT_GetConnectPacketSize_IgnoreAndReturn( MQTTSuccess );
+    incomingPacket.type = MQTT_PACKET_TYPE_CONNACK;
+    incomingPacket.remainingLength = 2;
+
+    /* Not enough time to receive entire packet, for branch coverage. */
+    timeout = 1;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTRecvFailed, status );
+
+    timeout = 10;
+
+    /* Not enough space for packet, discard it. */
+    mqttContext.networkBuffer.size = 2;
+    incomingPacket.remainingLength = 3;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTNoDataAvailable, status );
+
+    /* Timeout while discarding packet. */
+    incomingPacket.remainingLength = 20;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTRecvFailed, status );
+
+    /* Receive failure while discarding packet. */
+    mqttContext.transportInterface.recv = transportRecvFailure;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTRecvFailed, status );
+}
+
+/**
+ * @brief Test success case for MQTT_Connect().
+ */
+void test_MQTT_Connect_happy_path()
+{
+    MQTTContext_t mqttContext;
+    MQTTConnectInfo_t connectInfo;
+    MQTTPublishInfo_t willInfo;
+    uint32_t timeout = 2;
+    bool sessionPresent;
+    MQTTStatus_t status;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    MQTTPacketInfo_t incomingPacket;
+    size_t remainingLength, packetSize;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    MQTT_SerializeConnect_IgnoreAndReturn( MQTTSuccess );
+    MQTT_GetConnectPacketSize_IgnoreAndReturn( MQTTSuccess );
+
+    /* Success. */
+    incomingPacket.type = MQTT_PACKET_TYPE_CONNACK;
+    incomingPacket.remainingLength = 2;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    MQTT_DeserializeAck_IgnoreAndReturn( MQTTSuccess );
+    status = MQTT_Connect( &mqttContext, &connectInfo, NULL, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTSuccess, status );
+    TEST_ASSERT_EQUAL_INT( MQTTConnected, mqttContext.connectStatus );
+
+    /* With non-NULL Will. */
+    mqttContext.connectStatus = MQTTNotConnected;
+    willInfo.pTopicName = "test";
+    willInfo.topicNameLength = 4;
+    MQTT_GetIncomingPacketTypeAndLength_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetIncomingPacketTypeAndLength_ReturnThruPtr_pIncomingPacket( &incomingPacket );
+    status = MQTT_Connect( &mqttContext, &connectInfo, &willInfo, timeout, &sessionPresent );
+    TEST_ASSERT_EQUAL_INT( MQTTSuccess, status );
+    TEST_ASSERT_EQUAL_INT( MQTTConnected, mqttContext.connectStatus );
+}
+
+/* ========================================================================== */
+
+/**
+ * @brief Test that MQTT_Publish works as intended.
+ */
+void test_MQTT_Publish( void )
+{
+    MQTTContext_t mqttContext;
+    MQTTPublishInfo_t publishInfo;
+    uint16_t packetId;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    MQTTStatus_t status;
+    size_t headerSize;
+
+    const uint16_t PACKET_ID = 1;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+    transport.send = transportSendFailure;
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    memset( ( void * ) &publishInfo, 0x0, sizeof( publishInfo ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    /* Verify parameters. */
+    status = MQTT_Publish( NULL, &publishInfo, PACKET_ID );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+    status = MQTT_Publish( &mqttContext, NULL, PACKET_ID );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+    publishInfo.qos = MQTTQoS1;
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    /* Bad Parameter when getting packet size. */
+    publishInfo.qos = MQTTQoS0;
+    MQTT_GetPublishPacketSize_ExpectAnyArgsAndReturn( MQTTBadParameter );
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    /* Always return success from now on. */
+    MQTT_GetPublishPacketSize_IgnoreAndReturn( MQTTSuccess );
+
+    MQTT_SerializePublishHeader_ExpectAnyArgsAndReturn( MQTTNoMemory );
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTNoMemory, status );
+
+    /* The transport interface will fail. */
+    MQTT_SerializePublishHeader_ExpectAnyArgsAndReturn( MQTTSuccess );
+    /* We need sendPacket to be called with at least 1 byte to send, so that
+     * it can return failure. This argument is the output of serializing the
+     * publish header. */
+    headerSize = 1;
+    MQTT_SerializePublishHeader_ReturnThruPtr_pHeaderSize( &headerSize );
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTSendFailed, status );
+
+    /* We can ignore this now since MQTT_Publish initializes the header size to
+     * 0, so its initial send returns success (since 0 bytes are sent). */
+    MQTT_SerializePublishHeader_IgnoreAndReturn( MQTTSuccess );
+    publishInfo.pPayload = "Test";
+    publishInfo.payloadLength = 4;
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTSendFailed, status );
+
+    mqttContext.transportInterface.send = transportSendSuccess;
+    status = MQTT_Publish( &mqttContext, &publishInfo, 0 );
+    TEST_ASSERT_EQUAL_INT( MQTTSuccess, status );
+
+    /* Now for non zero QoS, which uses state engine. */
+    publishInfo.qos = MQTTQoS2;
+    MQTT_ReserveState_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_UpdateStatePublish_ExpectAnyArgsAndReturn( MQTTStateNull );
+    status = MQTT_Publish( &mqttContext, &publishInfo, PACKET_ID );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    publishInfo.qos = MQTTQoS1;
+    MQTT_ReserveState_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_UpdateStatePublish_ExpectAnyArgsAndReturn( MQTTPublishSend );
+    status = MQTT_Publish( &mqttContext, &publishInfo, PACKET_ID );
+    TEST_ASSERT_EQUAL_INT( MQTTSuccess, status );
+}
+
+/* ========================================================================== */
+
+/**
+ * @brief Test that MQTT_Disconnect works as intended.
+ */
+void test_MQTT_Disconnect( void )
+{
+    MQTTContext_t mqttContext;
+    MQTTStatus_t status;
+    uint8_t buffer[ 10 ];
+    uint8_t * bufPtr = buffer;
+    MQTTNetworkContext_t networkContext = ( MQTTNetworkContext_t ) &bufPtr;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    size_t disconnectSize = 2;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+    transport.networkContext = networkContext;
+    transport.recv = transportRecvSuccess;
+    transport.send = transportSendFailure;
+
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+    mqttContext.connectStatus = MQTTConnected;
+
+    /* Verify parameters. */
+    status = MQTT_Disconnect( NULL );
+    TEST_ASSERT_EQUAL_INT( MQTTBadParameter, status );
+
+    /* Send failure. */
+    MQTT_GetDisconnectPacketSize_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetDisconnectPacketSize_ReturnThruPtr_pPacketSize( &disconnectSize );
+    MQTT_SerializeDisconnect_ExpectAnyArgsAndReturn( MQTTSuccess );
+    status = MQTT_Disconnect( &mqttContext );
+    TEST_ASSERT_EQUAL( MQTTSendFailed, status );
+
+    /* Successful send. */
+    mqttContext.transportInterface.send = mockSend;
+    MQTT_GetDisconnectPacketSize_ExpectAnyArgsAndReturn( MQTTSuccess );
+    MQTT_GetDisconnectPacketSize_ReturnThruPtr_pPacketSize( &disconnectSize );
+    MQTT_SerializeDisconnect_ExpectAnyArgsAndReturn( MQTTSuccess );
+    /* Write a disconnect packet into the buffer. */
+    mqttBuffer[ 0 ] = MQTT_PACKET_TYPE_DISCONNECT;
+    status = MQTT_Disconnect( &mqttContext );
+    TEST_ASSERT_EQUAL( MQTTSuccess, status );
+    TEST_ASSERT_EQUAL( MQTTNotConnected, mqttContext.connectStatus );
+    TEST_ASSERT_EQUAL_MEMORY( mqttBuffer, buffer, 2 );
+}
+
+/* ========================================================================== */
+
+/**
+ * @brief Test that MQTT_GetPacketId works as intended.
+ */
+void test_MQTT_GetPacketId( void )
+{
+    MQTTContext_t mqttContext;
+    MQTTTransportInterface_t transport;
+    MQTTFixedBuffer_t networkBuffer;
+    MQTTApplicationCallbacks_t callbacks;
+    uint16_t packetId;
+
+    setupTransportInterface( &transport );
+    setupCallbacks( &callbacks );
+    setupNetworkBuffer( &networkBuffer );
+    memset( ( void * ) &mqttContext, 0x0, sizeof( mqttContext ) );
+    MQTT_Init( &mqttContext, &transport, &callbacks, &networkBuffer );
+
+    /* Verify parameters. */
+    packetId = MQTT_GetPacketId( NULL );
+    TEST_ASSERT_EQUAL_INT( 0, packetId );
+
+    packetId = MQTT_GetPacketId( &mqttContext );
+    TEST_ASSERT_EQUAL_INT( 1, packetId );
+    TEST_ASSERT_EQUAL_INT( 2, mqttContext.nextPacketId );
+
+    mqttContext.nextPacketId = UINT16_MAX;
+    packetId = MQTT_GetPacketId( &mqttContext );
+    TEST_ASSERT_EQUAL_INT( UINT16_MAX, packetId );
+    TEST_ASSERT_EQUAL_INT( 1, mqttContext.nextPacketId );
+}