Function test_large_tensor_transfer

cake-core/tests/protocol.rs:326–353

()

#[tokio::test]
async fn test_large_tensor_transfer() {
    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;

    // Realistic hidden state size: (1, 5120) x F16 = 10,240 bytes
    let tensor = make_f16_tensor(&[1, 5120]);
    let orig_bytes: Vec<u8> = tensor.data().to_vec();
    assert_eq!(orig_bytes.len(), 10240);

    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) => {
            assert_eq!(raw.data, orig_bytes);
        }
        other => panic!("expected Tensor, got {:?}", other),
    }

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

// ============================================================================
// Benchmarks (run with: cargo test --test protocol -- --ignored --nocapture)