Function test_batch_echo

cake-core/tests/protocol.rs:228–259

()

#[tokio::test]
async fn test_batch_echo() {
    let mock = MockWorker::bind().await;
    let addr = mock.addr();
    let worker = tokio::spawn(async move { mock.handle_one(None).await });

    let (mut stream, _) = connect_and_handshake(addr, None).await;

    let tensor = make_f16_tensor(&[1, 64]);
    let orig_bytes: Vec<u8> = tensor.data().to_vec();
    let batch = vec![
        ("model.layers.0".into(), 0usize, 0usize),
        ("model.layers.1".into(), 1, 1),
    ];

    Message::from_batch(&tensor, batch)
        .to_writer(&mut stream)
        .await
        .unwrap();

    let (_, resp) = Message::from_reader(&mut stream).await.unwrap();
    match resp {
        Message::Tensor(raw) => {
            let recovered = raw.to_tensor(&Device::Cpu).unwrap();
            assert_eq!(recovered.shape().dims(), &[1, 64]);
            assert_eq!(recovered.data().to_vec(), orig_bytes);
        }
        other => panic!("expected Tensor, got {:?}", other),
    }

    drop(stream);
    let _ = worker.await;
}

#[tokio::test]
async fn test_multiple_sequential_ops() {