Function test_goodbye_and_continue

cake-core/tests/protocol.rs:292–323

()

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

    // First op
    let tensor = make_f16_tensor(&[1, 64]);
    Message::single_op("model.layers.0", &tensor, 0, 0)
        .to_writer(&mut stream)
        .await
        .unwrap();
    let (_, resp) = Message::from_reader(&mut stream).await.unwrap();
    assert!(matches!(resp, Message::Tensor(_)));

    // Goodbye
    Message::Goodbye.to_writer(&mut stream).await.unwrap();
    let (_, resp) = Message::from_reader(&mut stream).await.unwrap();
    assert!(matches!(resp, Message::WorkerInfo(_)));

    // Another op after goodbye (worker should still be alive)
    Message::single_op("model.layers.0", &tensor, 1, 0)
        .to_writer(&mut stream)
        .await
        .unwrap();
    let (_, resp) = Message::from_reader(&mut stream).await.unwrap();
    assert!(matches!(resp, Message::Tensor(_)));

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

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