def test_client_download_model(watchman_service): """ Test client's ability to download the model """ client = Client(project=tu.GORDO_PROJECT, target=tu.GORDO_SINGLE_TARGET) models = client.download_model() assert isinstance(models, dict) assert isinstance(models[tu.GORDO_SINGLE_TARGET], BaseEstimator) # Can't download model for non-existent target with pytest.raises(ValueError): client = Client(project=tu.GORDO_PROJECT, target="non-existent-target") client.download_model()
def download_model(ctx: click.Context, output_dir: str): """ Download the actual model from the target and write to an output directory """ client = Client(*ctx.obj["args"], **ctx.obj["kwargs"]) models = client.download_model() # Iterate over mapping of models and save into their own sub dirs of the output_dir for target, model in models.items(): model_out_dir = os.path.join(output_dir, target) os.mkdir(model_out_dir) click.secho( f"Writing model '{target}' to directory: '{model_out_dir}'...", nl=False) serializer.dump(model, model_out_dir) click.secho(f"done") click.secho(f"Wrote all models to directory: {output_dir}", fg="green")