rust/bt-gatt/src/test_utils.rs

// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use futures::channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
use futures::future::{ready, Ready};
use futures::Stream;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::Arc;
use std::task::Poll;

use bt_common::{PeerId, Uuid};

use crate::central::ScanResult;
use crate::client::CharacteristicNotification;
use crate::{types::*, GattTypes};

#[derive(Default)]
struct FakePeerServiceInner {
    // Notifier that's used to send out notification.
    notifiers: HashMap<Handle, UnboundedSender<Result<CharacteristicNotification>>>,

    // Characteristics to return when `read_characteristic` and `discover_characteristics` are
    // called.
    characteristics: HashMap<Handle, (Characteristic, Vec<u8>)>,
}

#[derive(Clone)]
pub struct FakePeerService {
    inner: Arc<Mutex<FakePeerServiceInner>>,
}

impl FakePeerService {
    pub fn new() -> Self {
        Self { inner: Arc::new(Mutex::new(Default::default())) }
    }

    // Adds a characteristic so that it can be returned when discover/read method is
    // called.
    // Also triggers sending a characteristic value change notification to be sent.
    pub fn add_characteristic(&mut self, char: Characteristic, value: Vec<u8>) {
        let mut lock = self.inner.lock();
        let handle = char.handle;
        lock.characteristics.insert(handle, (char, value.clone()));
        if let Some(notifier) = lock.notifiers.get_mut(&handle) {
            notifier
                .unbounded_send(Ok(CharacteristicNotification {
                    handle,
                    value,
                    maybe_truncated: false,
                }))
                .expect("should succeed");
        }
    }

    // Sets expected characteristic value so that it can be used for validation when
    // write method is called.
    pub fn expect_characteristic_value(&mut self, handle: &Handle, value: Vec<u8>) {
        let mut lock = self.inner.lock();
        let Some(char) = lock.characteristics.get_mut(handle) else {
            panic!("Can't find characteristic {handle:?} to set expected value");
        };
        char.1 = value;
    }
}

impl crate::client::PeerService<FakeTypes> for FakePeerService {
    fn discover_characteristics(
        &self,
        uuid: Option<Uuid>,
    ) -> <FakeTypes as GattTypes>::CharacteristicDiscoveryFut {
        let lock = self.inner.lock();
        let mut result = Vec::new();
        for (_handle, (char, _value)) in &lock.characteristics {
            match uuid {
                Some(uuid) if uuid == char.uuid => result.push(char.clone()),
                None => result.push(char.clone()),
                _ => {}
            }
        }
        ready(Ok(result))
    }

    fn read_characteristic<'a>(
        &self,
        handle: &Handle,
        _offset: u16,
        buf: &'a mut [u8],
    ) -> <FakeTypes as GattTypes>::ReadFut<'a> {
        let read_characteristics = &(*self.inner.lock()).characteristics;
        let Some((_, value)) = read_characteristics.get(handle) else {
            return ready(Err(Error::Gatt(GattError::InvalidHandle)));
        };
        buf[..value.len()].copy_from_slice(value.as_slice());
        ready(Ok((value.len(), false)))
    }

    // For testing, should call `expect_characteristic_value` with the expected
    // value.
    fn write_characteristic<'a>(
        &self,
        handle: &Handle,
        _mode: WriteMode,
        _offset: u16,
        buf: &'a [u8],
    ) -> <FakeTypes as GattTypes>::WriteFut<'a> {
        let expected_characteristics = &(*self.inner.lock()).characteristics;
        // The write operation was not expected.
        let Some((_, expected)) = expected_characteristics.get(handle) else {
            panic!("Write operation to characteristic {handle:?} was not expected");
        };
        // Value written was not expected.
        if buf.len() != expected.len() || &buf[..expected.len()] != expected.as_slice() {
            panic!("Value written to characteristic {handle:?} was not expected: {buf:?}");
        }
        ready(Ok(()))
    }

    fn read_descriptor<'a>(
        &self,
        _handle: &Handle,
        _offset: u16,
        _buf: &'a mut [u8],
    ) -> <FakeTypes as GattTypes>::ReadFut<'a> {
        todo!()
    }

    fn write_descriptor<'a>(
        &self,
        _handle: &Handle,
        _offset: u16,
        _buf: &'a [u8],
    ) -> <FakeTypes as GattTypes>::WriteFut<'a> {
        todo!()
    }

    fn subscribe(&self, handle: &Handle) -> <FakeTypes as GattTypes>::NotificationStream {
        let (sender, receiver) = unbounded();
        (*self.inner.lock()).notifiers.insert(*handle, sender);
        receiver
    }
}

#[derive(Clone)]
pub struct FakeServiceHandle {
    pub uuid: Uuid,
    pub is_primary: bool,
    pub fake_service: FakePeerService,
}

impl crate::client::PeerServiceHandle<FakeTypes> for FakeServiceHandle {
    fn uuid(&self) -> Uuid {
        self.uuid
    }

    fn is_primary(&self) -> bool {
        self.is_primary
    }

    fn connect(&self) -> <FakeTypes as GattTypes>::ServiceConnectFut {
        futures::future::ready(Ok(self.fake_service.clone()))
    }
}

#[derive(Default)]
struct FakeClientInner {
    fake_services: Vec<FakeServiceHandle>,
}

#[derive(Clone)]
pub struct FakeClient {
    inner: Arc<Mutex<FakeClientInner>>,
}

impl FakeClient {
    pub fn new() -> Self {
        FakeClient { inner: Arc::new(Mutex::new(FakeClientInner::default())) }
    }

    /// Add a fake peer service to this client.
    pub fn add_service(&mut self, uuid: Uuid, is_primary: bool, fake_service: FakePeerService) {
        self.inner.lock().fake_services.push(FakeServiceHandle { uuid, is_primary, fake_service });
    }
}

impl crate::Client<FakeTypes> for FakeClient {
    fn peer_id(&self) -> PeerId {
        todo!()
    }

    fn find_service(&self, uuid: Uuid) -> <FakeTypes as GattTypes>::FindServicesFut {
        let fake_services = &self.inner.lock().fake_services;
        let mut filtered_services = Vec::new();
        for handle in fake_services {
            if handle.uuid == uuid {
                filtered_services.push(handle.clone());
            }
        }

        futures::future::ready(Ok(filtered_services))
    }
}

#[derive(Clone)]
pub struct ScannedResultStream {
    inner: Arc<Mutex<ScannedResultStreamInner>>,
}

#[derive(Default)]
pub struct ScannedResultStreamInner {
    result: Option<Result<ScanResult>>,
}

impl ScannedResultStream {
    pub fn new() -> Self {
        Self { inner: Arc::new(Mutex::new(ScannedResultStreamInner::default())) }
    }

    /// Set scanned result item to output from the stream.
    pub fn set_scanned_result(&mut self, item: Result<ScanResult>) {
        self.inner.lock().result = Some(item);
    }
}

impl Stream for ScannedResultStream {
    type Item = Result<ScanResult>;

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let mut lock = self.inner.lock();
        if lock.result.is_some() {
            Poll::Ready(lock.result.take())
        } else {
            // Never wake up, as if we never find another result
            Poll::Pending
        }
    }
}

pub struct FakeTypes {}

impl GattTypes for FakeTypes {
    type Central = FakeCentral;
    type ScanResultStream = ScannedResultStream;
    type Client = FakeClient;
    type ConnectFuture = Ready<Result<FakeClient>>;
    type PeerServiceHandle = FakeServiceHandle;
    type FindServicesFut = Ready<Result<Vec<FakeServiceHandle>>>;
    type PeerService = FakePeerService;
    type ServiceConnectFut = Ready<Result<FakePeerService>>;
    type CharacteristicDiscoveryFut = Ready<Result<Vec<Characteristic>>>;
    type NotificationStream = UnboundedReceiver<Result<CharacteristicNotification>>;
    type ReadFut<'a> = Ready<Result<(usize, bool)>>;
    type WriteFut<'a> = Ready<Result<()>>;
}

#[derive(Default)]
pub struct FakeCentralInner {
    clients: HashMap<PeerId, FakeClient>,
}

#[derive(Clone)]
pub struct FakeCentral {
    inner: Arc<Mutex<FakeCentralInner>>,
}

impl FakeCentral {
    pub fn new() -> Self {
        Self { inner: Arc::new(Mutex::new(FakeCentralInner::default())) }
    }

    pub fn add_client(&mut self, peer_id: PeerId, client: FakeClient) {
        let _ = self.inner.lock().clients.insert(peer_id, client);
    }
}

impl crate::Central<FakeTypes> for FakeCentral {
    fn scan(&self, _filters: &[crate::central::ScanFilter]) -> ScannedResultStream {
        ScannedResultStream::new()
    }

    fn connect(&self, peer_id: PeerId) -> <FakeTypes as GattTypes>::ConnectFuture {
        let clients = &self.inner.lock().clients;
        let res = match clients.get(&peer_id) {
            Some(client) => Ok(client.clone()),
            None => Err(Error::PeerDisconnected(peer_id)),
        };
        futures::future::ready(res)
    }
}