Function test_single_op_echo

cake-core/tests/protocol.rs:197–225

()

#[tokio::test]
async fn test_single_op_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, 128]);
    let orig_bytes: Vec<u8> = tensor.data().to_vec();

    Message::single_op("model.layers.0", &tensor, 0, 0)
        .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.dtype(), DType::F16);
            assert_eq!(recovered.shape().dims(), &[1, 128]);
            assert_eq!(recovered.data().to_vec(), orig_bytes);
        }
        other => panic!("expected Tensor, got {:?}", other),
    }

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

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