def test_set_model_from_file(mock_model, use_cluster):
try:
mock_model.create_torch_cnn(filepath="./torch_cnn.pt")
c = Client(None, use_cluster)
c.set_model_from_file("file_cnn", "./torch_cnn.pt", "TORCH", "CPU")
returned_model = c.get_model("file_cnn")
with open("./torch_cnn.pt", "rb") as f:
model = f.read()
assert model == returned_model
finally:
os.remove("torch_cnn.pt")
def test_set_model(mock_model, use_cluster):
model = mock_model.create_torch_cnn()
c = Client(None, use_cluster)
c.set_model("simple_cnn", model, "TORCH", "CPU")
returned_model = c.get_model("simple_cnn")
assert model == returned_model