update upstream tests

CrayLabs · Jul 19, 2024 · 6267515 · 6267515
1 parent 05329a1
commit 6267515
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/dragon/test_error_handling.py b/tests/dragon/test_error_handling.py
@@ -84,7 +84,7 @@ def setup_worker_manager_model_bytes(test_dir, monkeypatch: pytest.MonkeyPatch):
         comm_channel_type=FileSystemCommChannel,
     )
 
-    tensor_key = MessageHandler.build_tensor_key("key")
+    tensor_key = MessageHandler.build_tensor_key("key", feature_store.descriptor)
     model = MessageHandler.build_model(b"model", "model name", "v 0.0.1")
     request = MessageHandler.build_request(
         test_dir, model, [tensor_key], [tensor_key], [], None
@@ -116,7 +116,7 @@ def setup_worker_manager_model_key(test_dir, monkeypatch: pytest.MonkeyPatch):
         comm_channel_type=FileSystemCommChannel,
     )
 
-    tensor_key = MessageHandler.build_tensor_key("key")
+    tensor_key = MessageHandler.build_tensor_key("key", feature_store.descriptor)
     model_key = MessageHandler.build_model_key("model key")
     request = MessageHandler.build_request(
         test_dir, model_key, [tensor_key], [tensor_key], [], None