// This file is part of Astarte.
//
// Copyright 2024 - 2025 SECO Mind Srl
//
// Licensed 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.
//
// SPDX-License-Identifier: Apache-2.0

//! Client to send data to astarte, add interfaces or access properties.

use std::{future::Future, sync::Arc};

use astarte_interfaces::{interface::Retention, mapping::path::MappingPathError, MappingPath};
use chrono::{DateTime, Utc};
use tracing::{debug, trace, warn};

use crate::{
    aggregate::AstarteObject,
    error::{AggregationError, InterfaceTypeError},
    logging::security::{notify_security_event, SecurityEvent},
    retention::{
        memory::{ItemValue, VolatileItemError},
        Id, RetentionId, StoredRetention, StoredRetentionExt,
    },
    state::{SharedState, Status},
    store::StoreCapabilities,
    transport::{
        mqtt::{error::MqttError, Mqtt},
        Connection, Publish,
    },
    Timestamp,
};
use crate::{error::DynError, transport::Disconnect};
use crate::{
    event::DeviceEvent,
    store::wrapper::StoreWrapper,
    types::AstarteData,
    validate::{ValidatedIndividual, ValidatedObject},
    Error,
};

mod individual;
mod introspection;
mod object;
mod property;

/// Error generated by or received from the connection.
#[non_exhaustive]
#[derive(thiserror::Error, Debug)]
pub enum RecvError {
    /// A generic connection related error.
    ///
    /// Should be downcasted to access the underling specific connection error.
    #[error("connection error, {0:?}")]
    Connection(#[source] DynError),
    /// Couldn't parse the mapping path.
    #[error("invalid mapping path")]
    InvalidEndpoint(#[from] MappingPathError),
    /// Couldn't find an interface with the given name.
    #[error("couldn't find interface '{name}'")]
    InterfaceNotFound {
        /// Name of the missing interface.
        name: String,
    },
    /// Couldn't find missing mapping in the interface.
    #[error("couldn't find mapping {mapping} in interface {interface}")]
    MappingNotFound {
        /// Name of the interface.
        interface: String,
        /// Path of the missing mapping.
        mapping: String,
    },
    /// Received a message without timestamp, on an interface with `explicit_timestamp`
    #[error("message received with missing explicit timestamp on {interface_name}{path}")]
    MissingTimestamp {
        /// The interface we received the data on.
        interface_name: String,
        /// The mapping path we received on.
        path: String,
    },
    /// Received unset on property without `allow_unset`
    #[error("unset received on property {interface_name}{path} without allow_unset")]
    Unset {
        /// The interface we received the data on.
        interface_name: String,
        /// The mapping path we received on.
        path: String,
    },
    /// Invalid aggregation between the interface and the data.
    #[error("invalid aggregation between interface and data")]
    Aggregation(#[from] AggregationError),
    /// Invalid aggregation between the interface and the data.
    #[error("invalid interface type between interface and data")]
    InterfaceType(#[from] InterfaceTypeError),
    /// Received data on a device owned interface.
    #[error("received data on a device owned interface {interface}{path}")]
    Ownership {
        /// The interface we received on.
        interface: String,
        /// The endpoint we received on.
        path: String,
    },
    /// Error when the Device is disconnected from Astarte or client.
    ///
    /// This is an unrecoverable error for the SDK.
    #[error("disconnected from astarte")]
    Disconnected,
}

// Safe conversion for connection error
impl RecvError {
    pub(crate) fn mqtt_connection_error(value: MqttError) -> Self {
        RecvError::Connection(value.into())
    }

    #[cfg(feature = "message-hub")]
    pub(crate) fn grpc_connection_error(value: crate::transport::grpc::GrpcError) -> Self {
        RecvError::Connection(value.into())
    }
}

/// A trait representing the behavior of an Astarte device client.
///
/// A device client is responsible for interacting with the Astarte platform by sending properties
/// and datastreams, handling events, and managing device interfaces.
pub trait Client: Clone {
    /// Send an individual datastream on an interface.
    ///
    /// ```no_run
    /// use astarte_device_sdk::{
    ///     store::memory::MemoryStore, builder::DeviceBuilder,
    ///     transport::mqtt::MqttConfig, types::AstarteData, prelude::*,
    /// };
    /// use chrono::{TimeZone, Utc};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mqtt_config = MqttConfig::with_credential_secret("realm_id", "device_id", "credential_secret", "pairing_url");
    ///
    ///     let (mut client, connection) = DeviceBuilder::new().store(MemoryStore::new())
    ///         .connection(mqtt_config).build().await.unwrap();
    ///
    ///     let value: i32 = 42;
    ///     client.send_individual("my.interface.name", "/endpoint/path", value.into())
    ///         .await
    ///         .unwrap();
    /// }
    /// ```
    fn send_individual(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Send an individual datastream on an interface, with an explicit timestamp.
    ///
    /// ```no_run
    /// use astarte_device_sdk::{
    ///     store::memory::MemoryStore, builder::DeviceBuilder,
    ///     transport::mqtt::MqttConfig, types::AstarteData, prelude::*,
    /// };
    /// use chrono::{TimeZone, Utc};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mqtt_config = MqttConfig::with_credential_secret("realm_id", "device_id", "credential_secret", "pairing_url");
    ///
    ///     let (mut client, connection) = DeviceBuilder::new().store(MemoryStore::new())
    ///         .connection(mqtt_config).build().await.unwrap();
    ///
    ///     let value: i32 = 42;
    ///     let timestamp = Utc.timestamp_opt(1537449422, 0).unwrap();
    ///     client.send_individual_with_timestamp("my.interface.name", "/endpoint/path", value.into(), timestamp)
    ///         .await
    ///         .unwrap();
    /// }
    /// ```
    fn send_individual_with_timestamp(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
        timestamp: chrono::DateTime<chrono::Utc>,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Send an object datastream on an interface.
    ///
    /// The usage is the same of
    /// [`send_object_with_timestamp`](crate::Client::send_object_with_timestamp),
    /// without the timestamp.
    fn send_object(
        &mut self,
        interface_name: &str,
        base_path: &str,
        data: AstarteObject,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Send an object datastream on an interface, with an explicit timestamp.
    ///
    /// ```no_run
    /// use astarte_device_sdk::{
    ///     store::memory::MemoryStore, builder::DeviceBuilder,
    ///     transport::mqtt::MqttConfig, types::AstarteData, prelude::*,
    /// };
    /// # #[cfg(feature = "derive")]
    /// use astarte_device_sdk::IntoAstarteObject;
    /// # #[cfg(not(feature = "derive"))]
    /// # use astarte_device_sdk_derive::IntoAstarteObject;
    /// use chrono::{TimeZone, Utc};
    ///
    /// #[derive(IntoAstarteObject)]
    /// struct TestObject {
    ///     endpoint1: f64,
    ///     endpoint2: bool,
    /// }
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mqtt_config = MqttConfig::with_credential_secret("realm_id", "device_id", "credential_secret", "pairing_url");
    ///
    ///     let (mut client, connection) = DeviceBuilder::new().store(MemoryStore::new())
    ///         .connection(mqtt_config).build().await.unwrap();
    ///
    ///     let data = TestObject {
    ///         endpoint1: 1.34,
    ///         endpoint2: false
    ///     };
    ///     let timestamp = Utc.timestamp_opt(1537449422, 0).unwrap();
    ///     client.send_object_with_timestamp("my.interface.name", "/endpoint/path", data.try_into().unwrap(), timestamp)
    ///         .await
    ///         .unwrap();
    /// }
    /// ```
    fn send_object_with_timestamp(
        &mut self,
        interface_name: &str,
        base_path: &str,
        data: AstarteObject,
        timestamp: chrono::DateTime<chrono::Utc>,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Send an individual datastream on an interface.
    ///
    /// ```no_run
    /// use astarte_device_sdk::{
    ///     store::memory::MemoryStore, builder::DeviceBuilder,
    ///     transport::mqtt::MqttConfig, types::AstarteData, prelude::*,
    /// };
    /// use chrono::{TimeZone, Utc};
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mqtt_config = MqttConfig::with_credential_secret("realm_id", "device_id", "credential_secret", "pairing_url");
    ///
    ///     let (mut client, connection) = DeviceBuilder::new().store(MemoryStore::new())
    ///         .connection(mqtt_config).build().await.unwrap();
    ///
    ///     let value: i32 = 42;
    ///     client.set_property("my.interface.name", "/endpoint/path", value.into())
    ///         .await
    ///         .unwrap();
    /// }
    /// ```
    fn set_property(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Unset a device property.
    ///
    /// ```no_run
    /// use astarte_device_sdk::{
    ///     store::memory::MemoryStore, builder::DeviceBuilder,
    ///     transport::mqtt::MqttConfig, types::AstarteData, prelude::*,
    /// };
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let mqtt_config = MqttConfig::with_credential_secret("realm_id", "device_id", "credential_secret", "pairing_url");
    ///
    ///     let (mut device, _connection) = DeviceBuilder::new().store(MemoryStore::new())
    ///         .connection(mqtt_config).build().await.unwrap();
    ///
    ///     device
    ///         .unset_property("my.interface.name", "/endpoint/path",)
    ///         .await
    ///         .unwrap();
    /// }
    /// ```
    fn unset_property(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
    ) -> impl Future<Output = Result<(), Error>> + Send;

    /// Receives an event from Astarte.
    ///
    /// # Note
    ///
    /// An event can only be received once, so if the client is cloned only one of the clients
    /// instances will receive the message.
    fn recv(&self) -> impl Future<Output = Result<DeviceEvent, RecvError>> + Send;
}

/// A trait representing the behavior of an Astarte device client to disconnect itself from Astarte.
pub trait ClientDisconnect {
    /// Cleanly disconnects the client consuming it.
    fn disconnect(&mut self) -> impl Future<Output = Result<(), Error>> + Send;
}

/// Client to send and receive message to and form Astarte or access the Device properties.
///
/// ### Notes
///
/// Cloning the client will not broadcast the [`DeviceEvent`]. Each message can
/// only be received once.
#[derive(Debug)]
pub struct DeviceClient<C>
where
    C: Connection,
{
    // Sender of the connection.
    sender: C::Sender,
    // We use flume instead of the mpsc channel for the DeviceEvents for the connection to che
    // client since we need the Receiver end to be cloneable. Flume provides an async mpmc
    // channel/queue that fits our needs and doesn't suffer from the "slow receiver" problem.
    // Since it doesn't block the sender till all the receivers have read the msg. Unlike the
    // Tokio broadcast channel (another mpmc channel implementation).
    rx: flume::Receiver<Result<DeviceEvent, RecvError>>,
    pub(crate) store: StoreWrapper<C::Store>,
    pub(crate) state: Arc<SharedState>,
}

impl<C> DeviceClient<C>
where
    C: Connection,
{
    pub(crate) fn new(
        sender: C::Sender,
        rx: flume::Receiver<Result<DeviceEvent, RecvError>>,
        store: StoreWrapper<C::Store>,
        state: Arc<SharedState>,
    ) -> Self {
        Self {
            sender,
            rx,
            store,
            state,
        }
    }

    async fn send<T>(
        state: &SharedState,
        store: &StoreWrapper<C::Store>,
        sender: &mut C::Sender,
        data: T,
    ) -> Result<(), Error>
    where
        T: ClientPacket + TryInto<ItemValue, Error = VolatileItemError> + Clone,
        C::Store: StoreCapabilities,
        C::Sender: Publish,
    {
        match state.status.connection() {
            Status::Connected => {
                trace!("publish while connection is connected");
            }
            Status::Disconnected => {
                trace!("publish while connection is offline");
                return Self::offline_send(state, store, sender, data).await;
            }
            Status::Closed => {
                return Err(Error::Disconnected);
            }
        }

        match data.get_retention() {
            Retention::Volatile { .. } => Self::send_volatile(state, sender, data).await,
            Retention::Stored { .. } => Self::send_stored(state, store, sender, data).await,
            Retention::Discard => data.send(sender).await,
        }
    }

    async fn offline_send<T>(
        state: &SharedState,
        store: &StoreWrapper<C::Store>,
        sender: &mut C::Sender,
        data: T,
    ) -> Result<(), Error>
    where
        T: ClientPacket + TryInto<ItemValue, Error = VolatileItemError>,
        C::Store: StoreCapabilities,
        C::Sender: Publish,
    {
        match data.get_retention() {
            Retention::Discard => {
                debug!("drop publish with retention discard since disconnected");
            }
            Retention::Volatile { .. } => {
                let id = state.retention_ctx.next();

                state.volatile_store.push_unsent(id, data).await;
            }
            Retention::Stored { .. } => {
                let id = state.retention_ctx.next();

                if let Some(retention) = store.get_retention() {
                    data.store_publish(&id, sender, retention, false).await?;
                } else {
                    warn!(?store, "storing interface with retention 'Stored' in volatile store since the store doesn't support retention");
                    state.volatile_store.push_unsent(id, data).await;
                }
            }
        }

        Ok(())
    }

    async fn send_stored<T>(
        state: &SharedState,
        store: &StoreWrapper<C::Store>,
        sender: &mut C::Sender,
        data: T,
    ) -> Result<(), Error>
    where
        T: ClientPacket + TryInto<ItemValue, Error = VolatileItemError> + Clone,
        C::Store: StoreCapabilities,
        C::Sender: Publish,
    {
        let Some(retention) = store.get_retention() else {
            warn!(?store, "storing interface with retention 'Stored' in volatile store since the store doesn't support retention");
            return Self::send_volatile(state, sender, data).await;
        };

        // generate id after the check to avoid wasting an id generation in case it gets regenerated in send_volatile
        let id = state.retention_ctx.next();

        data.store_publish(&id, sender, retention, true).await?;
        data.send_stored(RetentionId::Stored(id), sender).await
    }

    async fn send_volatile<T>(
        state: &SharedState,
        sender: &mut C::Sender,
        data: T,
    ) -> Result<(), Error>
    where
        T: ClientPacket + TryInto<ItemValue, Error = VolatileItemError> + Clone,
        C::Store: StoreCapabilities,
        C::Sender: Publish,
    {
        let id = state.retention_ctx.next();

        state.volatile_store.push_sent(id, data.clone()).await;
        data.send_stored(RetentionId::Volatile(id), sender).await
    }
}

impl<S> DeviceClient<Mqtt<S>>
where
    S: StoreCapabilities,
{
    /// Retrieve the expiry (not_after) timestamp of the current certificate
    pub async fn get_cert_expiry(&self) -> Option<DateTime<Utc>> {
        *self.state.cert_expiry.lock().await
    }

    /// Retrieve the expiry (not_after) timestamp of the current certificate
    /// Note that this function will log a security event if the feature is enabled
    /// when the certificate will expire at the passed datetime
    pub async fn is_valid_at(&self, check_dt: DateTime<Utc>) -> Option<bool> {
        let expiry = self.get_cert_expiry().await?;

        if check_dt < expiry {
            Some(true)
        } else {
            notify_security_event(SecurityEvent::CertificateAboutToExpire);

            Some(false)
        }
    }
}

// Cannot be derived it has specific generic bounds.
impl<C> Clone for DeviceClient<C>
where
    C: Connection,
{
    fn clone(&self) -> Self {
        Self {
            sender: self.sender.clone(),
            rx: self.rx.clone(),
            store: self.store.clone(),
            state: Arc::clone(&self.state),
        }
    }
}

impl<C> Client for DeviceClient<C>
where
    C: Connection,
    C::Sender: Publish,
{
    async fn send_object_with_timestamp(
        &mut self,
        interface_name: &str,
        base_path: &str,
        data: AstarteObject,
        timestamp: chrono::DateTime<chrono::Utc>,
    ) -> Result<(), Error> {
        let path = MappingPath::try_from(base_path)?;

        self.send_datastream_object(interface_name, &path, data, Some(timestamp))
            .await
    }

    async fn send_object(
        &mut self,
        interface_name: &str,
        base_path: &str,
        data: AstarteObject,
    ) -> Result<(), Error> {
        let path = MappingPath::try_from(base_path)?;

        self.send_datastream_object(interface_name, &path, data, None)
            .await
    }

    async fn send_individual(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
    ) -> Result<(), Error> {
        let path = MappingPath::try_from(mapping_path)?;

        self.send_datastream_individual(interface_name, &path, data, None)
            .await
    }

    async fn send_individual_with_timestamp(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
        timestamp: chrono::DateTime<chrono::Utc>,
    ) -> Result<(), Error> {
        let mapping = MappingPath::try_from(mapping_path)?;

        self.send_datastream_individual(interface_name, &mapping, data, Some(timestamp))
            .await
    }

    async fn set_property(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
        data: AstarteData,
    ) -> Result<(), Error> {
        trace!("setting property {}{}", interface_name, mapping_path);

        let path = MappingPath::try_from(mapping_path)?;

        self.send_property(interface_name, &path, data).await
    }

    async fn unset_property(
        &mut self,
        interface_name: &str,
        mapping_path: &str,
    ) -> Result<(), Error> {
        trace!("unsetting {}{}", interface_name, mapping_path);

        let path = MappingPath::try_from(mapping_path)?;

        self.send_unset(interface_name, &path).await
    }

    async fn recv(&self) -> Result<DeviceEvent, RecvError> {
        self.rx
            .recv_async()
            .await
            .map_err(|_| RecvError::Disconnected)?
    }
}

impl<C> ClientDisconnect for DeviceClient<C>
where
    C: Connection,
    C::Sender: Disconnect,
{
    async fn disconnect(&mut self) -> Result<(), Error> {
        self.sender.disconnect().await?;

        self.state.status.close();

        Ok(())
    }
}

trait ClientPacket {
    fn get_retention(&self) -> Retention;

    fn serialize<S>(&self, sender: &S) -> Result<Vec<u8>, crate::Error>
    where
        S: Publish;

    fn send<S>(self, sender: &mut S) -> impl Future<Output = Result<(), crate::Error>> + Send
    where
        S: Publish + Send;

    fn send_stored<S>(
        self,
        id: RetentionId,
        sender: &mut S,
    ) -> impl Future<Output = Result<(), crate::Error>> + Send
    where
        S: Publish + Send;

    fn store_publish<S, R>(
        &self,
        id: &Id,
        sender: &S,
        retention: &R,
        sent: bool,
    ) -> impl Future<Output = Result<(), crate::Error>> + Send
    where
        S: Publish + Sync,
        R: StoredRetention + Sync;
}

impl ClientPacket for ValidatedIndividual {
    fn get_retention(&self) -> Retention {
        self.retention
    }

    fn serialize<S>(&self, sender: &S) -> Result<Vec<u8>, crate::Error>
    where
        S: Publish,
    {
        sender.serialize_individual(self)
    }

    async fn send<S>(self, sender: &mut S) -> Result<(), crate::Error>
    where
        S: Publish + Send,
    {
        sender.send_individual(self).await
    }

    async fn send_stored<S>(self, id: RetentionId, sender: &mut S) -> Result<(), crate::Error>
    where
        S: Publish,
    {
        sender.send_individual_stored(id, self).await
    }

    async fn store_publish<S, R>(
        &self,
        id: &Id,
        sender: &S,
        retention: &R,
        sent: bool,
    ) -> Result<(), crate::Error>
    where
        S: Publish + Sync,
        R: StoredRetention + Sync,
    {
        let serialized = self.serialize(sender)?;

        retention
            .store_publish_individual(id, self, &serialized, sent)
            .await
            .map_err(crate::Error::from)
    }
}

impl ClientPacket for ValidatedObject {
    fn get_retention(&self) -> Retention {
        self.retention
    }

    fn serialize<S>(&self, sender: &S) -> Result<Vec<u8>, crate::Error>
    where
        S: Publish,
    {
        sender.serialize_object(self)
    }

    async fn send<S>(self, sender: &mut S) -> Result<(), crate::Error>
    where
        S: Publish + Send,
    {
        sender.send_object(self).await
    }

    async fn send_stored<S>(self, id: RetentionId, sender: &mut S) -> Result<(), crate::Error>
    where
        S: Publish + Send,
    {
        sender.send_object_stored(id, self).await
    }

    async fn store_publish<S, R>(
        &self,
        id: &Id,
        sender: &S,
        retention: &R,
        sent: bool,
    ) -> Result<(), crate::Error>
    where
        S: Publish + Sync,
        R: StoredRetention + Sync,
    {
        let serialized = self.serialize(sender)?;

        retention
            .store_publish_object(id, self, &serialized, sent)
            .await
            .map_err(crate::Error::from)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use std::str::FromStr;

    use astarte_interfaces::Interface;
    use chrono::Utc;
    use mockall::Sequence;
    use pretty_assertions::assert_eq;

    use crate::builder::{DEFAULT_CHANNEL_SIZE, DEFAULT_VOLATILE_CAPACITY};
    use crate::interfaces::Interfaces;
    use crate::retention::memory::VolatileStore;
    use crate::state::Status;
    use crate::store::memory::MemoryStore;
    use crate::store::StoreCapabilities;
    use crate::transport::mock::{MockCon, MockSender};
    use crate::Value;

    use super::*;

    pub(crate) fn mock_client(
        interfaces: &[&str],
    ) -> (
        DeviceClient<MockCon<MemoryStore>>,
        flume::Sender<Result<DeviceEvent, RecvError>>,
    ) {
        mock_client_with_store(interfaces, MemoryStore::new())
    }

    pub(crate) fn mock_client_with_store<S>(
        interfaces: &[&str],
        store: S,
    ) -> (
        DeviceClient<MockCon<S>>,
        flume::Sender<Result<DeviceEvent, RecvError>>,
    )
    where
        S: StoreCapabilities,
    {
        let interfaces = interfaces.iter().map(|i| Interface::from_str(i).unwrap());
        let interfaces = Interfaces::from_iter(interfaces);

        let sender = MockSender::new();
        let (tx, rx) = flume::bounded(DEFAULT_CHANNEL_SIZE);
        let state = SharedState::new(
            interfaces,
            VolatileStore::with_capacity(DEFAULT_VOLATILE_CAPACITY),
        );

        let client = DeviceClient::new(sender, rx, StoreWrapper::new(store), Arc::new(state));

        (client, tx)
    }

    #[test]
    fn client_must_be_clone() {
        let (mut client, _tx) = mock_client(&[]);

        let mut seq = Sequence::new();
        client
            .sender
            .expect_clone()
            .once()
            .in_sequence(&mut seq)
            .returning(MockSender::new);

        let _b = client.clone();
    }

    #[tokio::test]
    async fn client_recv() {
        let (client, tx) = mock_client(&[]);

        let exp = DeviceEvent {
            interface: "interface".to_string(),
            path: "path".to_string(),
            data: Value::Individual {
                data: AstarteData::LongInteger(42),
                timestamp: Utc::now(),
            },
        };

        tx.send_async(Ok(exp.clone())).await.unwrap();

        let event = client.recv().await.unwrap();

        assert_eq!(event, exp);
    }

    #[tokio::test]
    async fn client_disconnect_closed() {
        let (mut client, _tx) = mock_client(&[]);

        let mut seq = Sequence::new();
        client
            .sender
            .expect_disconnect()
            .once()
            .in_sequence(&mut seq)
            .returning(|| Ok(()));

        client.disconnect().await.unwrap();

        assert_eq!(client.state.status.connection(), Status::Closed);
    }
}
