fix tests

src/client/dusty_block.py

@@ -42,7 +42,7 @@ class DustyRemoteBlock(RemoteExpert):
                     is_meta_set = True
                 return chunk
 
-            return rpc_call(amap_in_executor(_metadata_setter, input), timeout)
+            return await rpc_call(amap_in_executor(_metadata_setter, input), timeout)
 
         return rpc
 

tests/test_dust_payment.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from typing import AsyncIterator, Optional
 
 import pytest
@@ -28,8 +26,7 @@ class HandlerStubTest:
         return input
 
     async def rpc_stream(self, input: AsyncIterator[ExpertRequest], timeout: Optional[float] = None):
-        async for i in input:
-            yield i
+        return input
 
     async def rpc_info(self, input: str, timeout: Optional[float] = None):
         return input
@@ -48,6 +45,7 @@ async def test_single():
     input = torch.randn(1, 2)
     request = ExpertRequest(uid="expert1", tensors=[serialize_torch_tensor(input)])
 
+    print(stub)
     out: ExpertRequest = await stub.rpc_single(request)
 
     assert out.metadata != b""
@@ -87,3 +85,12 @@ async def test_stream():
     results = await deserialize_tensor_stream(amap_in_executor(lambda r: r.tensors, iter_as_aiter(outputs_list)))
     assert len(results) == 1
     assert torch.allclose(results[0], input)
+
+
+@pytest.mark.asyncio
+async def test_no_wrapper():
+    remote = RemoteBlockTest(DustBankTest(), None, None)
+    stub = remote.stub
+
+    test = await stub.rpc_info("Test")
+    assert test == "Test"