mirror of
https://github.com/chanzuckerberg/cellxgene.git
synced 2026-09-15 20:57:56 +08:00
ddb601c103
* Replace --input-h5ad-file with a positional argument, for consistency with other CLI commands * Replace --update-h5ad-file with --overwrite, for consistency with `prepare` command. * Fix/clarify various help descriptions * Fix final output message when input file is overwritten * Fix annotate top-level help description
131 lines
5.1 KiB
Python
131 lines
5.1 KiB
Python
import os
|
|
import shutil
|
|
import unittest
|
|
from tempfile import mkstemp, TemporaryDirectory, NamedTemporaryFile
|
|
|
|
import mlflow
|
|
from click.testing import CliRunner
|
|
|
|
from server.cli.annotate import annotate
|
|
from test.unit.cli.fixtures.mlflow_model_fixture import FakeModel
|
|
|
|
|
|
def write_model(model) -> str:
|
|
with TemporaryDirectory() as mlflow_model_dir:
|
|
fixtures_path = os.path.join(os.path.dirname(__file__), "fixtures")
|
|
mlflow.pyfunc.save_model(mlflow_model_dir, loader_module="fixtures", code_path=[fixtures_path])
|
|
return shutil.make_archive(mkstemp()[1], "zip", mlflow_model_dir)
|
|
|
|
|
|
class TestCliAnnotate(unittest.TestCase):
|
|
def test__annotate__loads_and_runs(self):
|
|
"""
|
|
Invokes the `annotate` subcommand of cellxgene CLI, using a CliRunner() programmatic invocation.
|
|
|
|
This tests the happy path case:
|
|
1) Command line options are parsed;
|
|
2) An MLflow model zip archive can be read in (from local disk), unpacked, and invoked;
|
|
3) The correct options are passed to the MLflow model.
|
|
4) The annotate subcommand exits successfully.
|
|
|
|
This does not verify model output or predictions (it's a fake MLflow model, after all); it's up to the real model
|
|
to output its predictions as it wants, but this is specific to the model and so not tested here.
|
|
|
|
The CliRunner() invokes the subcommand in a subprocess, and the annotate subcommand itself invokes the MLflow
|
|
model in yet another subprocess. So while this test can help determine if everything is working, it is not a
|
|
simple matter to debug in the case of a failure. However, the stdout/stderr of the MLflow process is captured
|
|
by the CliRunner() subprocess, so errors can be inspected in result.stdout when debugging this test. Hope this
|
|
helps!
|
|
"""
|
|
|
|
_, query_dataset_file_path = mkstemp()
|
|
model_file_path = write_model(FakeModel())
|
|
|
|
result = CliRunner().invoke(
|
|
annotate,
|
|
[
|
|
query_dataset_file_path,
|
|
"--model-url",
|
|
model_file_path,
|
|
"--output-h5ad-file",
|
|
f"{query_dataset_file_path}.output",
|
|
# avoid having mflow create conda env or virtualenv when in test env;
|
|
# this avoids making pip remote requests and is also faster
|
|
"--mlflow-env-manager",
|
|
"local",
|
|
],
|
|
)
|
|
|
|
# to help debugging, show the output from the CliRunner and MLflow stdout
|
|
if result.exit_code:
|
|
print(result.stdout)
|
|
|
|
self.assertEqual(0, result.exit_code, "runs successfully")
|
|
|
|
# The FakeModel will print it inputs to stdout, as "__MODEL_INPUT__={...}", allowing us to assert that it received valid inputs.
|
|
self.assertIn(
|
|
"__MODEL_INPUT__={"
|
|
f'"query_dataset_h5ad_path": "{query_dataset_file_path}", '
|
|
f'"output_h5ad_path": "{query_dataset_file_path}.output", '
|
|
'"annotation_prefix": "cxg_cell_type", "classifier": "default", '
|
|
'"organism": "Homo sapiens", "use_gpu": true}',
|
|
result.stdout,
|
|
"inputs passed correctly",
|
|
)
|
|
self.assertIn(
|
|
f"Wrote annotations to {query_dataset_file_path}.output",
|
|
result.stdout,
|
|
"success message is correct",
|
|
)
|
|
|
|
def test__annotate__requires_overwrite_option_when_output_file_exists(self):
|
|
|
|
with NamedTemporaryFile() as input_h5ad, NamedTemporaryFile() as existing_file:
|
|
required_options = [input_h5ad.name, "--output-h5ad-file", existing_file.name, "--model-url", "some_url"]
|
|
result = CliRunner().invoke(
|
|
annotate,
|
|
required_options + [],
|
|
)
|
|
|
|
self.assertNotEqual(0, result.exit_code, "aborts with non-success code")
|
|
self.assertIn(
|
|
"try using the flag --overwrite",
|
|
result.stdout,
|
|
"error message displayed",
|
|
)
|
|
|
|
def test__annotate__overwrite_option_allows_overwrite_of_existing_output_file(self):
|
|
model_file_path = write_model(FakeModel())
|
|
|
|
with NamedTemporaryFile() as existing_file:
|
|
required_options = [
|
|
existing_file.name,
|
|
"--output-h5ad-file",
|
|
existing_file.name,
|
|
"--overwrite",
|
|
"--model-url",
|
|
model_file_path,
|
|
]
|
|
result = CliRunner().invoke(
|
|
annotate,
|
|
required_options + [],
|
|
)
|
|
|
|
print(result.stdout)
|
|
self.assertNotEqual(1, result.exit_code, "aborts with non-success code")
|
|
self.assertIn(
|
|
f"Wrote annotations to {existing_file.name}",
|
|
result.stdout,
|
|
"success message is correct on output file overwrite",
|
|
)
|
|
|
|
|
|
# TODO:
|
|
# Test annotate cli args more comprehensively
|
|
# Test server.cli.annotate._validate_options
|
|
# Test model caching feature works
|
|
# Test model loading from s3 works (maybe w/just a real model)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|