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::{AdvertisingDatum, PeerName, 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
    }
}

pub struct FakeServiceHandle {}

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

    fn is_primary(&self) -> bool {
        todo!()
    }

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

pub struct FakeClient {}

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

    fn find_service(&self, _uuid: Uuid) -> <FakeTypes as GattTypes>::FindServicesFut {
        futures::future::ready(Ok::<Vec<FakeServiceHandle>, Error>(vec![FakeServiceHandle {}]))
    }
}

pub struct SingleResultStream {
    result: Option<Result<crate::central::ScanResult>>,
}

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

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        _cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        if self.result.is_some() {
            Poll::Ready(self.get_mut().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 = SingleResultStream;
    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 FakeCentral {}

impl crate::Central<FakeTypes> for FakeCentral {
    fn scan(&self, _filters: &[crate::central::ScanFilter]) -> SingleResultStream {
        SingleResultStream {
            result: Some(Ok(ScanResult {
                id: PeerId(1),
                connectable: true,
                name: PeerName::CompleteName("Marie's Pixel 7 Pro".to_owned()),
                advertised: vec![AdvertisingDatum::Services(vec![Uuid::from_u16(0x1844)])],
            })),
        }
    }

    fn connect(&self, _peer_id: PeerId) -> <FakeTypes as GattTypes>::ConnectFuture {
        futures::future::ready(Ok(FakeClient {}))
    }
}