uwb(uci-rust): Cleanup MockUwbAdaption

Clients using mocks should ideally be only dealing with the interface
exposed by the corresponding module. The current impl of
MockUwbAdaptation needs clients to dig into the internals of
UwbClientCallback.

For ex: Uci module tests have to deal with packet fragmentation inside
UwbClientCallback impl which should be hidden from unit tests of
clients.

Bug: 216552887
Test: atest libuwb_uci_rust_tests
Change-Id: I071daedd7b1f92089ac4db28707c44baf7a1598e

3 files changed, 85 insertions, 83 deletions

diff --git a/src/fuzz/fuzzer.rs b/src/fuzz/fuzzer.rs
index f56739f..72688c7 100644
--- a/src/fuzz/fuzzer.rs
+++ b/src/fuzz/fuzzer.rs
@@ -6,12 +6,12 @@ use log::{error, info};
 use num_traits::cast::FromPrimitive;
 use std::sync::Arc;
 use tokio::sync::mpsc;
-use uwb_uci_packets::{SessionInitRspBuilder, StatusCode};
+use uwb_uci_packets::{SessionInitRspBuilder, StatusCode, UciCommandPacket};
 use uwb_uci_rust::adaptation::MockUwbAdaptation;
 use uwb_uci_rust::error::UwbErr;
 use uwb_uci_rust::event_manager::MockEventManager;
 use uwb_uci_rust::uci::{
-    uci_hmsgs, Dispatcher, DispatcherImpl, HalCallback, JNICommand, SyncUwbAdaptation,
+    uci_hmsgs, uci_hrcv, Dispatcher, DispatcherImpl, HalCallback, JNICommand, SyncUwbAdaptation,
 };
 
 fn create_dispatcher_with_mock_adaptation(msgs: Vec<JNICommand>) -> Result<DispatcherImpl, UwbErr> {
@@ -40,11 +40,15 @@ fn create_dispatcher_with_mock_adaptation(msgs: Vec<JNICommand>) -> Result<Dispa
     )
 }
 
-fn generate_fake_cmd_rsp(msg: &JNICommand) -> Result<(Vec<u8>, Vec<u8>), UwbErr> {
+fn generate_fake_cmd_rsp(
+    msg: &JNICommand,
+) -> Result<(UciCommandPacket, uci_hrcv::UciResponse), UwbErr> {
     match msg {
         JNICommand::UciSessionInit(session_id, session_type) => Ok((
             uci_hmsgs::build_session_init_cmd(*session_id, *session_type)?.build().into(),
-            SessionInitRspBuilder { status: StatusCode::UciStatusOk }.build().into(),
+            uci_hrcv::UciResponse::SessionInitRsp(
+                SessionInitRspBuilder { status: StatusCode::UciStatusOk }.build(),
+            ),
         )),
         _ => Err(UwbErr::Undefined),
     }
diff --git a/src/rust/adaptation/mod.rs b/src/rust/adaptation/mod.rs
index 5f6f1c3..5e1c8fd 100644
--- a/src/rust/adaptation/mod.rs
+++ b/src/rust/adaptation/mod.rs
@@ -2,7 +2,6 @@
 
 use crate::error::UwbErr;
 use crate::uci::uci_hrcv;
-use crate::uci::uci_logger::MockUciLogger;
 use crate::uci::uci_logger::{UciLogMode, UciLogger, UciLoggerImpl};
 use crate::uci::HalCallback;
 use android_hardware_uwb::aidl::android::hardware::uwb::{
@@ -216,9 +215,9 @@ enum ExpectedCall {
         out: Result<()>,
     },
     SendUciMessage {
-        expected_data: Vec<u8>,
-        rsp_data: Option<Vec<u8>>,
-        notf_data: Option<Vec<u8>>,
+        expected_cmd: UciCommandPacket,
+        rsp: Option<uci_hrcv::UciResponse>,
+        notf: Option<uci_hrcv::UciNotification>,
         out: Result<()>,
     },
 }
@@ -262,36 +261,29 @@ impl MockUwbAdaptation {
     #[allow(dead_code)]
     pub fn expect_send_uci_message(
         &self,
-        expected_data: Vec<u8>,
-        rsp_data: Option<Vec<u8>>,
-        notf_data: Option<Vec<u8>>,
+        expected_cmd: UciCommandPacket,
+        rsp: Option<uci_hrcv::UciResponse>,
+        notf: Option<uci_hrcv::UciNotification>,
         out: Result<()>,
     ) {
         self.expected_calls.lock().unwrap().push_back(ExpectedCall::SendUciMessage {
-            expected_data,
-            rsp_data,
-            notf_data,
+            expected_cmd,
+            rsp,
+            notf,
             out,
         });
     }
 
-    fn create_uwb_client_callback(
-        rsp_sender: mpsc::UnboundedSender<HalCallback>,
-    ) -> UwbClientCallback {
-        // Add tests for the mock logger.
-        UwbClientCallback::new(rsp_sender, Arc::new(MockUciLogger::new()))
+    async fn send_hal_event(&self, event: UwbEvent, event_status: UwbStatus) {
+        self.rsp_sender.send(HalCallback::Event { event, event_status }).unwrap();
     }
 
-    async fn send_client_event(&self, event: UwbEvent, status: UwbStatus) {
-        let uwb_client_callback =
-            MockUwbAdaptation::create_uwb_client_callback(self.rsp_sender.clone());
-        let _ = uwb_client_callback.onHalEvent(event, status).await;
+    async fn send_uci_response(&self, rsp: uci_hrcv::UciResponse) {
+        self.rsp_sender.send(HalCallback::UciRsp(rsp)).unwrap();
     }
 
-    async fn send_client_message(&self, rsp_data: Vec<u8>) {
-        let uwb_client_callback =
-            MockUwbAdaptation::create_uwb_client_callback(self.rsp_sender.clone());
-        let _ = uwb_client_callback.onUciMessage(&rsp_data).await;
+    async fn send_uci_notification(&self, ntf: uci_hrcv::UciNotification) {
+        self.rsp_sender.send(HalCallback::UciNtf(ntf)).unwrap();
     }
 }
 
@@ -335,7 +327,7 @@ impl UwbAdaptation for MockUwbAdaptation {
         match expected_out {
             Some(out) => {
                 let status = if out.is_ok() { UwbStatus::OK } else { UwbStatus::FAILED };
-                self.send_client_event(UwbEvent::OPEN_CPLT, status).await;
+                self.send_hal_event(UwbEvent::OPEN_CPLT, status).await;
                 out
             }
             None => {
@@ -361,7 +353,7 @@ impl UwbAdaptation for MockUwbAdaptation {
         match expected_out {
             Some(out) => {
                 let status = if out.is_ok() { UwbStatus::OK } else { UwbStatus::FAILED };
-                self.send_client_event(UwbEvent::CLOSE_CPLT, status).await;
+                self.send_hal_event(UwbEvent::CLOSE_CPLT, status).await;
                 out
             }
             None => {
@@ -387,7 +379,7 @@ impl UwbAdaptation for MockUwbAdaptation {
         match expected_out {
             Some(out) => {
                 let status = if out.is_ok() { UwbStatus::OK } else { UwbStatus::FAILED };
-                self.send_client_event(UwbEvent::POST_INIT_CPLT, status).await;
+                self.send_hal_event(UwbEvent::POST_INIT_CPLT, status).await;
                 out
             }
             None => {
@@ -427,11 +419,13 @@ impl UwbAdaptation for MockUwbAdaptation {
         let expected_out = {
             let mut expected_calls = self.expected_calls.lock().unwrap();
             match expected_calls.pop_front() {
-                Some(ExpectedCall::SendUciMessage { expected_data, rsp_data, notf_data, out })
-                    if expected_data == cmd.to_vec() =>
-                {
-                    Some((rsp_data, notf_data, out))
-                }
+                Some(ExpectedCall::SendUciMessage {
+                    expected_cmd,
+                    rsp,
+                    notf,
+                    out,
+                    // PDL generated packets do not implement PartialEq, so use the raw bytes for comparison.
+                }) if expected_cmd.clone().to_bytes() == cmd.to_bytes() => Some((rsp, notf, out)),
                 Some(call) => {
                     expected_calls.push_front(call);
                     None
@@ -441,12 +435,12 @@ impl UwbAdaptation for MockUwbAdaptation {
         };
 
         match expected_out {
-            Some((rsp_data, notf_data, out)) => {
-                if let Some(rsp) = rsp_data {
-                    self.send_client_message(rsp).await;
+            Some((rsp, notf, out)) => {
+                if let Some(rsp) = rsp {
+                    self.send_uci_response(rsp).await;
                 }
-                if let Some(notf) = notf_data {
-                    self.send_client_message(notf).await;
+                if let Some(notf) = notf {
+                    self.send_uci_notification(notf).await;
                 }
                 out
             }
@@ -459,12 +453,6 @@ impl UwbAdaptation for MockUwbAdaptation {
 }
 
 #[cfg(test)]
-fn create_uwb_client_callback(rsp_sender: mpsc::UnboundedSender<HalCallback>) -> UwbClientCallback {
-    // Add tests for the mock logger.
-    UwbClientCallback::new(rsp_sender, Arc::new(MockUciLogger::new()))
-}
-
-#[cfg(test)]
 pub mod tests {
     #![allow(non_snake_case)]
     use super::*;
@@ -479,6 +467,15 @@ pub mod tests {
         SessionInit { expected_session_id: i32, out: BinderResult<()> },
         SendUciMessage { expected_data: Vec<u8>, out: BinderResult<i32> },
     }
+    use crate::uci::uci_logger::MockUciLogger;
+
+    #[cfg(test)]
+    fn create_uwb_client_callback(
+        rsp_sender: mpsc::UnboundedSender<HalCallback>,
+    ) -> UwbClientCallback {
+        // Add tests for the mock logger.
+        UwbClientCallback::new(rsp_sender, Arc::new(MockUciLogger::new()))
+    }
 
     pub struct MockHal {
         expected_calls: StdMutex<VecDeque<ExpectedHalCall>>,
diff --git a/src/rust/uci/mod.rs b/src/rust/uci/mod.rs
index a24d2f2..a8d8893 100644
--- a/src/rust/uci/mod.rs
+++ b/src/rust/uci/mod.rs
@@ -556,16 +556,16 @@ impl Dispatcher for DispatcherImpl {
 
 #[cfg(test)]
 mod tests {
+    use self::uci_hrcv::UciNotification;
+    use self::uci_hrcv::UciResponse;
+
     use super::*;
     use crate::adaptation::MockUwbAdaptation;
     use crate::event_manager::MockEventManager;
     use android_hardware_uwb::aidl::android::hardware::uwb::{
         UwbEvent::UwbEvent, UwbStatus::UwbStatus,
     };
-    use uwb_uci_packets::{
-        DeviceState, DeviceStatusNtfBuilder, GetDeviceInfoRspBuilder, Packet, UciPacketHalPacket,
-        UciPacketPacket,
-    };
+    use uwb_uci_packets::*;
 
     fn setup_dispatcher(
         config_fn: fn(&mut Arc<MockUwbAdaptation>, &mut MockEventManager),
@@ -609,9 +609,9 @@ mod tests {
         Ok((dispatcher, rsp_sender))
     }
 
-    fn generate_fake_cmd_rsp_data() -> (Vec<u8>, Vec<u8>) {
-        let cmd_data = GetDeviceInfoCmdBuilder {}.build().to_vec();
-        let rsp_packet: UciPacketPacket = GetDeviceInfoRspBuilder {
+    fn generate_fake_get_device_cmd_rsp() -> (GetDeviceInfoCmdPacket, GetDeviceInfoRspPacket) {
+        let cmd = GetDeviceInfoCmdBuilder {}.build();
+        let rsp = GetDeviceInfoRspBuilder {
             status: StatusCode::UciStatusOk,
             uci_version: 0,
             mac_version: 0,
@@ -619,21 +619,12 @@ mod tests {
             uci_test_version: 0,
             vendor_spec_info: vec![],
         }
-        .build()
-        .into();
-        // Convert to UciPacketHalPacket
-        let mut rsp_frags: Vec<UciPacketHalPacket> = rsp_packet.into();
-        let rsp_data = rsp_frags.pop().unwrap().to_vec();
-
-        (cmd_data, rsp_data)
+        .build();
+        (cmd, rsp)
     }
 
-    fn generate_fake_ntf_data() -> Vec<u8> {
-        let ntf_packet: UciPacketPacket =
-            DeviceStatusNtfBuilder { device_state: DeviceState::DeviceStateReady }.build().into();
-        // Convert to UciPacketHalPacket
-        let mut ntf_frags: Vec<UciPacketHalPacket> = ntf_packet.into();
-        ntf_frags.pop().unwrap().to_vec()
+    fn generate_fake_device_status_ntf() -> DeviceStatusNtfPacket {
+        DeviceStatusNtfBuilder { device_state: DeviceState::DeviceStateReady }.build()
     }
 
     #[test]
@@ -674,10 +665,15 @@ mod tests {
     }
 
     #[test]
-    fn test_send_uci_message() -> Result<()> {
+    fn test_get_device_info() -> Result<()> {
         let mut dispatcher = setup_dispatcher(|mock_adaptation, _mock_event_manager| {
-            let (cmd_data, rsp_data) = generate_fake_cmd_rsp_data();
-            mock_adaptation.expect_send_uci_message(cmd_data, Some(rsp_data), None, Ok(()));
+            let (cmd, rsp) = generate_fake_get_device_cmd_rsp();
+            mock_adaptation.expect_send_uci_message(
+                cmd.into(),
+                Some(UciResponse::GetDeviceInfoRsp(rsp)),
+                None,
+                Ok(()),
+            );
         })?;
 
         dispatcher.block_on_jni_command(JNICommand::UciGetDeviceInfo)?;
@@ -685,14 +681,19 @@ mod tests {
     }
 
     #[test]
-    fn test_send_uci_message_with_retry() -> Result<()> {
+    fn test_get_device_info_with_uci_retry() -> Result<()> {
         let mut dispatcher = setup_dispatcher(|mock_adaptation, _mock_event_manager| {
-            let (cmd_data, rsp_data) = generate_fake_cmd_rsp_data();
+            let (cmd, rsp) = generate_fake_get_device_cmd_rsp();
 
             // Let the first 2 tries not response data, then the 3rd tries response successfully.
-            mock_adaptation.expect_send_uci_message(cmd_data.clone(), None, None, Ok(()));
-            mock_adaptation.expect_send_uci_message(cmd_data.clone(), None, None, Ok(()));
-            mock_adaptation.expect_send_uci_message(cmd_data, Some(rsp_data), None, Ok(()));
+            mock_adaptation.expect_send_uci_message(cmd.clone().into(), None, None, Ok(()));
+            mock_adaptation.expect_send_uci_message(cmd.clone().into(), None, None, Ok(()));
+            mock_adaptation.expect_send_uci_message(
+                cmd.into(),
+                Some(UciResponse::GetDeviceInfoRsp(rsp)),
+                None,
+                Ok(()),
+            );
         })?;
 
         dispatcher.block_on_jni_command(JNICommand::UciGetDeviceInfo)?;
@@ -700,10 +701,10 @@ mod tests {
     }
 
     #[test]
-    fn test_send_uci_message_failed() -> Result<()> {
+    fn test_get_device_info_send_uci_message_failed() -> Result<()> {
         let dispatcher = setup_dispatcher(|mock_adaptation, _mock_event_manager| {
-            let (cmd_data, _rsp_data) = generate_fake_cmd_rsp_data();
-            mock_adaptation.expect_send_uci_message(cmd_data, None, None, Err(UwbErr::failed()));
+            let (cmd, _rsp) = generate_fake_get_device_cmd_rsp();
+            mock_adaptation.expect_send_uci_message(cmd.into(), None, None, Err(UwbErr::failed()));
         })?;
 
         dispatcher
@@ -713,14 +714,14 @@ mod tests {
     }
 
     #[test]
-    fn test_notification() -> Result<()> {
+    fn test_device_status_notification() -> Result<()> {
         let mut dispatcher = setup_dispatcher(|mock_adaptation, mock_event_manager| {
-            let (cmd_data, rsp_data) = generate_fake_cmd_rsp_data();
-            let ntf_data = generate_fake_ntf_data();
+            let (cmd, rsp) = generate_fake_get_device_cmd_rsp();
+            let ntf = generate_fake_device_status_ntf();
             mock_adaptation.expect_send_uci_message(
-                cmd_data,
-                Some(rsp_data),
-                Some(ntf_data),
+                cmd.into(),
+                Some(UciResponse::GetDeviceInfoRsp(rsp)),
+                Some(UciNotification::DeviceStatusNtf(ntf)),
                 Ok(()),
             );
             mock_event_manager.expect_device_status_notification_received(Ok(()));