Add user-generated annotations tests to the server (#1164)

* Add user-generated annotations tests to the server

Partially completes https://github.com/chanzuckerberg/cellxgene/issues/969

* Auto-format python code

* @skip_if: passing lambdas > than property strings

* Respond to feedback from @bkmartinjr

client/Makefile

@@ -1,7 +1,7 @@
 include ../common.mk
 
 DATASET := $(if $(DATASET),$(DATASET),../example-dataset/pbmc3k.h5ad)
-ANNOTATIONS := $(if $(ANNOTATIONS),$(ANNOTATIONS),../example-dataset/pbmc3k-annotations.csv)
+ANNOTATIONS := $(if $(ANNOTATIONS),$(ANNOTATIONS),../server/test/test_datasets/pbmc3k-annotations.csv)
 ANNOTATIONS_FILENAME := $(shell basename $(ANNOTATIONS))
 
 # Packaging

pyproject.toml

@@ -17,7 +17,7 @@ exclude = '''
     | buck-out
     | build
     | dist
-    | server/app/util/fbs/NetEncoding
+    | server/data_common/fbs/NetEncoding
   )/
 
 )

server/app/app.py

@@ -82,12 +82,12 @@ def rest_get_data_adaptor(func):
 
 def static_redirect(dataset, therest):
     """ redirect all static requests to the standard location """
-    return redirect(f'/static/{therest}', code=301)
+    return redirect(f"/static/{therest}", code=301)
 
 
 def favicon_redirect(dataset):
     """ redirect favicon to static dir """
-    return redirect('/static/favicon.png', code=301)
+    return redirect("/static/favicon.png", code=301)
 
 
 def dataroot_index():
@@ -127,20 +127,17 @@ class SchemaAPI(Resource):
 class ConfigAPI(Resource):
     @rest_get_data_adaptor
     def get(self, data_adaptor):
-        return common_rest.config_get(
-            current_app.app_config, data_adaptor, current_app.annotations)
+        return common_rest.config_get(current_app.app_config, data_adaptor, current_app.annotations)
 
 
 class AnnotationsObsAPI(Resource):
     @rest_get_data_adaptor
     def get(self, data_adaptor):
-        return common_rest.annotations_obs_get(
-            request, data_adaptor, current_app.annotations)
+        return common_rest.annotations_obs_get(request, data_adaptor, current_app.annotations)
 
     @rest_get_data_adaptor
     def put(self, data_adaptor):
-        return common_rest.annotations_obs_put(
-            request, data_adaptor, current_app.annotations)
+        return common_rest.annotations_obs_put(request, data_adaptor, current_app.annotations)
 
 
 class AnnotationsVarAPI(Resource):
@@ -216,7 +213,7 @@ class Server:
             bp_api = Blueprint("api_dataset", __name__, url_prefix="/<dataset>" + api_version)
             resources = get_api_resources(bp_api)
             self.app.register_blueprint(resources.blueprint)
-            self.app.add_url_rule("/<dataset>/", 'dataset_index', dataset_index)
+            self.app.add_url_rule("/<dataset>/", "dataset_index", dataset_index)
             self.app.add_url_rule("/<dataset>/static/<path:therest>", "static_redirect", static_redirect)
             self.app.add_url_rule("/<dataset>/favicon.png", "favicon_redirect", favicon_redirect)
 

server/cli/launch.py

@@ -20,7 +20,7 @@ from server.common.errors import OntologyLoadFailure
 
 # anything bigger than this will generate a special message
 BIG_FILE_SIZE_THRESHOLD = 100 * 2 ** 20  # 100MB
-DEFAULT_SERVER_PORT = int(environ.get('CXG_SERVER_PORT', '5005'))
+DEFAULT_SERVER_PORT = int(environ.get("CXG_SERVER_PORT", "5005"))
 
 
 def annotation_args(func):
@@ -54,14 +54,14 @@ def annotation_args(func):
         is_flag=True,
         default=False,
         show_default=True,
-        help="When creating annotations, optionally autocomplete names from ontology terms."
+        help="When creating annotations, optionally autocomplete names from ontology terms.",
     )
     @click.option(
         "--experimental-annotations-ontology-obo",
         default=None,
         show_default=True,
         metavar="<path or url>",
-        help="Location of OBO file defining cell annotation autosuggest terms."
+        help="Location of OBO file defining cell annotation autosuggest terms.",
     )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
@@ -110,7 +110,6 @@ def config_args(func):
 
 
 def dataset_args(func):
-
     @click.option(
         "--obs-names",
         "-obs",
@@ -131,15 +130,9 @@ def dataset_args(func):
         is_flag=True,
         default=False,
         show_default=False,
-        help="Load anndata in file-backed mode. "
-             "This may save memory, but may result in slower overall performance.",
-    )
-    @click.option(
-        "--title",
-        "-t",
-        metavar="<text>",
-        help="Title to display. If omitted will use file name."
+        help="Load anndata in file-backed mode. " "This may save memory, but may result in slower overall performance.",
     )
+    @click.option("--title", "-t", metavar="<text>", help="Title to display. If omitted will use file name.")
     @click.option(
         "--about",
         metavar="<URL>",
@@ -212,8 +205,9 @@ def launch_args(func):
         default=None,
         metavar="<data directory>",
         help="Enable cellxgene to serve multiple files. Supply path (local directory or URL)"
-             " to folder containing H5AD and/or CXG datasets.",
-        hidden=True)  # TODO, unhide when dataroot is supported)
+        " to folder containing H5AD and/or CXG datasets.",
+        hidden=True,
+    )  # TODO, unhide when dataroot is supported)
     @click.argument("datapath", required=False, metavar="<path to data file>")
     @click.option(
         "--open",
@@ -282,7 +276,7 @@ def launch(
     backed,
     disable_diffexp,
     experimental_annotations_ontology,
-    experimental_annotations_ontology_obo
+    experimental_annotations_ontology_obo,
 ):
     """Launch the cellxgene data viewer.
     This web app lets you explore single-cell expression data.
@@ -308,7 +302,7 @@ def launch(
 
     if datapath is None and dataroot is None:
         # TODO:  change the error message once dataroot is fully supported
-        raise click.ClickException("Missing argument \"<path to data file>.\"")
+        raise click.ClickException('Missing argument "<path to data file>."')
         # raise click.ClickException("must supply either <path to data file> or --dataroot")
     if datapath is not None and dataroot is not None:
         raise click.ClickException("must supply only one of <path to data file> or --dataroot")
@@ -376,6 +370,7 @@ def launch(
                 )
 
     if about:
+
         def url_check(url):
             try:
                 result = urlparse(url)
@@ -405,7 +400,8 @@ def launch(
         obs_names=obs_names,
         var_names=var_names,
         anndata_backed=backed,
-        disable_diffexp=disable_diffexp)
+        disable_diffexp=disable_diffexp,
+    )
 
     matrix_data_cache_manager = MatrixDataCacheManager()
     data_adaptor = None
@@ -424,8 +420,7 @@ def launch(
     annotations = None
 
     if experimental_annotations:
-        annotations = AnnotationsLocalFile(experimental_annotations_output_dir,
-                                           experimental_annotations_file)
+        annotations = AnnotationsLocalFile(experimental_annotations_output_dir, experimental_annotations_file)
 
         # if the user has specified a fixed label file, go ahead and validate it
         # so that we can remove errors early in the process.
@@ -441,6 +436,7 @@ def launch(
 
     # create the server
     from server.app.app import Server
+
     server = Server(matrix_data_cache_manager, annotations, app_config)
 
     if not verbose:

server/cli/upgrade.py

@@ -20,7 +20,7 @@ def log_upgrade_check():
 
     # Get the current latest release
     try:
-        release_tag_generator = (r['tag_name'] for r in _request_cellxgene_releases())
+        release_tag_generator = (r["tag_name"] for r in _request_cellxgene_releases())
         latest_release = next(release_tag_generator, lambda tag_name: validate_version_str(tag_name))
         if version_gt(latest_release, __version__):
             click.echo(f"There's a new version of cellxgene available ({latest_release})!")
@@ -37,15 +37,16 @@ class RateLimitException(Exception):
 
 def _request_cellxgene_releases():
     def raise_on_rate_limit(response):
-        if response.status_code == 403 and res.headers.get('X-RateLimit-Remaining') == '0':
+        if response.status_code == 403 and res.headers.get("X-RateLimit-Remaining") == "0":
             raise RateLimitException
+
     url = "https://api.github.com/repos/chanzuckerberg/cellxgene/releases"
     res = requests.get(url)
     raise_on_rate_limit(res)
     for release in res.json():
         yield release
-    while 'next' in res.links.keys():
-        res = requests.get(res.links['next']['url'])
+    while "next" in res.links.keys():
+        res = requests.get(res.links["next"]["url"])
         raise_on_rate_limit(res)
         for release in res.json():
             yield release

server/common/annotations.py

@@ -42,7 +42,7 @@ class Annotations(metaclass=ABCMeta):
             raise OntologyLoadFailure(f"Unable to find OBO ontology path: {path}") from e
 
         except SyntaxError as e:
-            msg = ''.join(traceback.format_exception_only(SyntaxError, e))
+            msg = "".join(traceback.format_exception_only(SyntaxError, e))
             raise OntologyLoadFailure(msg) from e
 
         except Exception as e:

server/common/app_config.py

@@ -13,16 +13,12 @@ class AppFeature(object):
             setattr(self, k, v)
 
     def todict(self):
-        d = dict(
-            available=self.available,
-            method=self.method,
-            path=self.path)
+        d = dict(available=self.available, method=self.method, path=self.path)
         d.update(self.extra)
         return d
 
 
 class AppConfig(object):
-
     def __init__(self, **kw):
         super().__init__()
 
@@ -47,10 +43,20 @@ class AppConfig(object):
         # parameters
         self.diffexp_may_be_slow = False
 
-        inputs = ["datapath", "dataroot", "title", "about", "scripts", "layout",
-                  "max_category_items", "diffexp_lfc_cutoff",
-                  "obs_names", "var_names",
-                  "anndata_backed", "disable_diffexp"]
+        inputs = [
+            "datapath",
+            "dataroot",
+            "title",
+            "about",
+            "scripts",
+            "layout",
+            "max_category_items",
+            "diffexp_lfc_cutoff",
+            "obs_names",
+            "var_names",
+            "anndata_backed",
+            "disable_diffexp",
+        ]
 
         self.update(inputs, kw)
 
@@ -80,9 +86,7 @@ class AppConfig(object):
         title = self.get_title(data_adaptor)
         about = self.get_about(data_adaptor)
 
-        display_names = dict(
-            engine=data_adaptor.get_name(),
-            dataset=title)
+        display_names = dict(engine=data_adaptor.get_name(), dataset=title)
 
         # library_versions
         library_versions = {}
@@ -90,7 +94,7 @@ class AppConfig(object):
         library_versions["cellxgene"] = cellxgene_version
 
         # links
-        links = {"about-dataset" : about}
+        links = {"about-dataset": about}
 
         # parameters
         parameters = {

server/common/data_locator.py

@@ -84,7 +84,7 @@ class DataLocator:
         # and clean it up when done.  If the path has a suffix/extension,
         # do our best to create a file with the same.
         ext = os.path.splitext(self.path)
-        suffix = None if ext[1] == '' else ext[1]
+        suffix = None if ext[1] == "" else ext[1]
         with self.open() as src, tempfile.NamedTemporaryFile(prefix="cellxgene_", suffix=suffix, delete=False) as tmp:
             tmp.write(src.read())
             tmp.close()

server/common/errors.py

@@ -2,6 +2,7 @@ class FilterError(Exception):
     """
     Raised when filter is malformed
     """
+
     pass
 
 
@@ -9,6 +10,7 @@ class JSONEncodingValueError(Exception):
     """
     Raised when data cannot be encoded into json
     """
+
     pass
 
 
@@ -16,6 +18,7 @@ class MimeTypeError(Exception):
     """
     Raised when incompatible MIME type selected
     """
+
     pass
 
 
@@ -23,6 +26,7 @@ class PrepareError(Exception):
     """
     Raised when data is misprepared
     """
+
     pass
 
 
@@ -30,6 +34,7 @@ class DatasetAccessError(Exception):
     """
     Raised when file loaded into a DataAdaptor is misformatted
     """
+
     pass
 
 
@@ -37,6 +42,7 @@ class DisabledFeatureError(Exception):
     """
     Raised when an attempt to use a disabled feature occurs
     """
+
     pass
 
 
@@ -44,6 +50,7 @@ class AnnotationsError(Exception):
     """
     Raised when an attempt to use the annotations feature fails
     """
+
     pass
 
 
@@ -51,4 +58,5 @@ class OntologyLoadFailure(Exception):
     """
     Raised when reading the ontology file fails
     """
+
     pass

server/common/rest.py

@@ -29,9 +29,7 @@ def schema_get_helper(data_adaptor, annotations):
 
 def schema_get(data_adaptor, annotations):
     schema = schema_get_helper(data_adaptor, annotations)
-    return make_response(
-        jsonify({"schema": schema}), HTTPStatus.OK
-    )
+    return make_response(jsonify({"schema": schema}), HTTPStatus.OK)
 
 
 def config_get(app_config, data_adaptor, annotations):

server/common/utils.py

@@ -92,20 +92,18 @@ def jsonify_numpy(data):
 def dtype_to_schema(dtype):
     schema = {}
     if dtype == np.float32:
-        schema['type'] = 'float32'
+        schema["type"] = "float32"
     elif dtype == np.int32:
-        schema['type'] = 'int32'
+        schema["type"] = "int32"
     elif dtype == np.bool_:
-        schema['type'] = 'boolean'
+        schema["type"] = "boolean"
     elif dtype == np.str:
-        schema['type'] = 'string'
+        schema["type"] = "string"
     elif dtype == "category":
         schema["type"] = "categorical"
         schema["categories"] = dtype.categories.tolist()
     else:
-        raise TypeError(
-            f"Annotations of type {dtype} are unsupported."
-        )
+        raise TypeError(f"Annotations of type {dtype} are unsupported.")
     return schema
 
 

server/converters/cxgtool.py

@@ -287,7 +287,7 @@ def create_emb(e_name, emb):
     * large tile size (1000)
     * default compression level
     """
-    filters = tiledb.FilterList([tiledb.ZstdFilter(), ])
+    filters = tiledb.FilterList([tiledb.ZstdFilter()])
     attrs = [tiledb.Attr(dtype=emb.dtype, filters=filters)]
     dims = []
     for d in range(emb.ndim):

server/data_anndata/anndata_adaptor.py

@@ -23,7 +23,6 @@ def anndata_version_is_pre_070():
 
 
 class AnndataAdaptor(DataAdaptor):
-
     def __init__(self, data_locator, config=None):
         super().__init__(config)
         self.data = None
@@ -127,7 +126,6 @@ class AnndataAdaptor(DataAdaptor):
             "dataframe": {"nObs": self.cell_count, "nVar": self.gene_count, "type": str(self.data.X.dtype)},
             "annotations": {
                 "obs": {"index": self.parameters.get("obs_names"), "columns": []},
-
                 "var": {"index": self.parameters.get("var_names"), "columns": []},
             },
             "layout": {"obs": []},
@@ -177,8 +175,10 @@ class AnndataAdaptor(DataAdaptor):
 
     def _validate_and_initialize(self):
         if anndata_version_is_pre_070() and self.config.anndata_backed:
-            warnings.warn(f"Use of --backed mode with anndata versions older than 0.7 will have serious "
-                          "performance issues. Please update to at least anndata 0.7 or later.")
+            warnings.warn(
+                f"Use of --backed mode with anndata versions older than 0.7 will have serious "
+                "performance issues. Please update to at least anndata 0.7 or later."
+            )
 
         # var and obs column names must be unique
         if not self.data.obs.columns.is_unique or not self.data.var.columns.is_unique:

server/data_common/data_adaptor.py

@@ -132,16 +132,14 @@ class DataAdaptor(metaclass=ABCMeta):
 
         if self.get_embedding_names():
             # TODO handle "var" when gene layout becomes available
-            features["layout_obs"] = AppFeature(
-                "/layout/obs", available=True)
+            features["layout_obs"] = AppFeature("/layout/obs", available=True)
         else:
             features["layout_obs"] = AppFeature("/layout/obs")
 
         if self.config.disable_diffexp:
             features["diffexp"] = AppFeature("/diffexp/")
         else:
-            features["diffexp"] = AppFeature(
-                "/diffexp/", available=True)
+            features["diffexp"] = AppFeature("/diffexp/", available=True)
 
         return features
 
@@ -152,17 +150,17 @@ class DataAdaptor(metaclass=ABCMeta):
         mask = np.zeros((count,), dtype=np.bool)
         for i in filter:
             if type(i) == list:
-                mask[i[0]: i[1]] = True
+                mask[i[0] : i[1]] = True
             else:
                 mask[i] = True
         return mask
 
     def _axis_filter_to_mask(self, axis, filter, count):
-        mask = np.ones((count, ), dtype=np.bool)
-        if 'index' in filter:
-            mask = np.logical_and(mask, self._index_filter_to_mask(filter['index'], count))
-        if 'annotation_value' in filter:
-            mask = np.logical_and(mask, self._annotation_filter_to_mask(axis, filter['annotation_value'], count))
+        mask = np.ones((count,), dtype=np.bool)
+        if "index" in filter:
+            mask = np.logical_and(mask, self._index_filter_to_mask(filter["index"], count))
+        if "annotation_value" in filter:
+            mask = np.logical_and(mask, self._annotation_filter_to_mask(axis, filter["annotation_value"], count))
 
         return mask
 
@@ -176,7 +174,7 @@ class DataAdaptor(metaclass=ABCMeta):
                 anno_data = self.query_obs_array(name)
 
             if anno_data.dtype.name in ["boolean", "category", "object"]:
-                values = v.get('values', [])
+                values = v.get("values", [])
                 key_idx = np.in1d(anno_data, values)
                 mask = np.logical_and(mask, key_idx)
 
@@ -202,10 +200,10 @@ class DataAdaptor(metaclass=ABCMeta):
         obs_selector = None
         if filter is not None:
             if Axis.OBS in filter:
-                obs_selector = self._axis_filter_to_mask(Axis.OBS, filter['obs'], shape[0])
+                obs_selector = self._axis_filter_to_mask(Axis.OBS, filter["obs"], shape[0])
 
             if Axis.VAR in filter:
-                var_selector = self._axis_filter_to_mask(Axis.VAR, filter['var'], shape[1])
+                var_selector = self._axis_filter_to_mask(Axis.VAR, filter["var"], shape[1])
 
         return (obs_selector, var_selector)
 
@@ -309,7 +307,7 @@ class DataAdaptor(metaclass=ABCMeta):
 
         embeddings = self.get_embedding_names()
         layout_data = []
-        with ServerTiming.time(f'layout.query'):
+        with ServerTiming.time(f"layout.query"):
             for ename in embeddings:
                 embedding = self.get_embedding_array(ename, 2)
 
@@ -326,7 +324,7 @@ class DataAdaptor(metaclass=ABCMeta):
                 normalized_layout = normalized_layout.astype(dtype=np.float32)
                 layout_data.append(pd.DataFrame(normalized_layout, columns=[f"{ename}_0", f"{ename}_1"]))
 
-        with ServerTiming.time(f'layout.encode'):
+        with ServerTiming.time(f"layout.encode"):
             if layout_data:
                 df = pd.concat(layout_data, axis=1, copy=False)
             else:

server/data_common/matrix_loader.py

@@ -135,7 +135,6 @@ class MatrixDataType(Enum):
 
 
 class MatrixDataLoader(object):
-
     def __init__(self, location, etype=None):
         self.location = location
         if etype is None:
@@ -145,9 +144,11 @@ class MatrixDataLoader(object):
         self.matrix_type = None
         if self.etype == MatrixDataType.H5AD:
             from server.data_anndata.anndata_adaptor import AnndataAdaptor
+
             self.matrix_type = AnndataAdaptor
         elif self.etype == MatrixDataType.CXG:
             from server.data_cxg.cxg_adaptor import CxgAdaptor
+
             self.matrix_type = CxgAdaptor
 
     def matrix_data_type(self):

server/data_common/rwlock.py

@@ -25,6 +25,7 @@ from threading import Lock
 # _______________________________________________________________________
 # Class
 
+
 class RWLock(object):
     """ RWLock class; this is meant to allow an object to be read from by
         multiple threads, but only written to by a single thread at a time. See:

server/data_cxg/cxg_adaptor.py

@@ -16,10 +16,7 @@ import threading
 class CxgAdaptor(DataAdaptor):
 
     # TODO:  The tiledb context parameters should be a configuration option
-    tiledb_ctx = tiledb.Ctx({
-        'sm.tile_cache_size': 8 * 1024 * 1024 * 1024,
-        'sm.num_reader_threads': 32,
-    })
+    tiledb_ctx = tiledb.Ctx({"sm.tile_cache_size": 8 * 1024 * 1024 * 1024, "sm.num_reader_threads": 32})
 
     def __init__(self, location, config=None):
         super().__init__(config)
@@ -28,8 +25,8 @@ class CxgAdaptor(DataAdaptor):
         self.lock = threading.Lock()
 
         self.url = location
-        if self.url[-1] != '/':
-            self.url += '/'
+        if self.url[-1] != "/":
+            self.url += "/"
 
         self._validate_and_initialize()
 
@@ -79,19 +76,18 @@ class CxgAdaptor(DataAdaptor):
         returns list of (absolute paths, type) *without* trailing slash
         in the path.
         """
+
         def _cleanpath(p):
-            if p[-1] == '/':
+            if p[-1] == "/":
                 return p[:-1]
             else:
                 return p
 
-        if uri[-1] != '/':
-            uri += '/'
+        if uri[-1] != "/":
+            uri += "/"
 
         result = []
-        tiledb.ls(uri,
-                  lambda path, type: result.append((_cleanpath(path), type)),
-                  ctx=self.tiledb_ctx)
+        tiledb.ls(uri, lambda path, type: result.append((_cleanpath(path), type)), ctx=self.tiledb_ctx)
         return result
 
     @staticmethod
@@ -135,13 +131,13 @@ class CxgAdaptor(DataAdaptor):
         elif a_type == "array":
             # version >0
             gmd = self.open_array("cxg_group_metadata")
-            cxg_version = gmd.meta['cxg_version']
+            cxg_version = gmd.meta["cxg_version"]
             if cxg_version == "0.1":
-                cxg_properties = json.loads(gmd.meta['cxg_properties'])
-                title = cxg_properties.get('title', None)
-                about = cxg_properties.get('about', None)
+                cxg_properties = json.loads(gmd.meta["cxg_properties"])
+                title = cxg_properties.get("title", None)
+                about = cxg_properties.get("about", None)
 
-        if cxg_version not in ['0.0', '0.1']:
+        if cxg_version not in ["0.0", "0.1"]:
             raise DatasetAccessError(f"cxg matrix is not valid: {self.url}")
 
         self.title = title
@@ -175,7 +171,7 @@ class CxgAdaptor(DataAdaptor):
         if obs_items == slice(None) and var_items == slice(None):
             data = X[:, :]
         else:
-            data = X.multi_index[obs_items, var_items]['']
+            data = X.multi_index[obs_items, var_items][""]
         return data
 
     def get_shape(self):
@@ -222,11 +218,9 @@ class CxgAdaptor(DataAdaptor):
     # function to get the embedding
     # this function to iterate through embeddings.
     def get_embedding_names(self):
-        with ServerTiming.time(f'layout.lsuri'):
+        with ServerTiming.time(f"layout.lsuri"):
             pemb = self.get_path("emb")
-            embeddings = [
-                os.path.basename(p) for (p, t) in self.lsuri(pemb) if t == 'array'
-            ]
+            embeddings = [os.path.basename(p) for (p, t) in self.lsuri(pemb) if t == "array"]
         return embeddings
 
     @staticmethod
@@ -235,82 +229,68 @@ class CxgAdaptor(DataAdaptor):
         dtype = attr.dtype
         schema = {}
         # type hints take precedence
-        if 'type' in type_hint:
-            schema['type'] = type_hint['type']
+        if "type" in type_hint:
+            schema["type"] = type_hint["type"]
         elif dtype == np.float32:
-            schema['type'] = 'float32'
+            schema["type"] = "float32"
         elif dtype == np.int32:
-            schema['type'] = 'int32'
+            schema["type"] = "int32"
         elif dtype == np.bool_:
-            schema['type'] = 'boolean'
+            schema["type"] = "boolean"
         elif dtype == np.str:
-            schema['type'] = 'string'
+            schema["type"] = "string"
         elif dtype == "category":
             schema["type"] = "categorical"
             schema["categories"] = dtype.categories.tolist()
         else:
-            raise TypeError(
-                f"Annotations of type {dtype} are unsupported."
-            )
+            raise TypeError(f"Annotations of type {dtype} are unsupported.")
 
-        if schema['type'] == 'categorical' and 'categories' in schema_hints:
-            schema['categories'] = schema_hints['categories']
+        if schema["type"] == "categorical" and "categories" in schema_hints:
+            schema["categories"] = schema_hints["categories"]
         return schema
 
     def get_schema(self):
         shape = self.get_shape()
         dtype = self.get_X_array_dtype()
 
-        dataframe = {
-            'nObs': shape[0],
-            'nVar': shape[1],
-            'type': dtype.name
-        }
+        dataframe = {"nObs": shape[0], "nVar": shape[1], "type": dtype.name}
 
         annotations = {}
-        for ax in ('obs', 'var'):
+        for ax in ("obs", "var"):
             A = self.open_array(ax)
-            schema_hints = json.loads(A.meta['cxg_schema']) if 'cxg_schema' in A.meta else {}
+            schema_hints = json.loads(A.meta["cxg_schema"]) if "cxg_schema" in A.meta else {}
             if type(schema_hints) is not dict:
-                raise TypeError(f'Array schema was malformed.')
+                raise TypeError(f"Array schema was malformed.")
 
             cols = []
             for attr in A.schema:
                 schema = dict(name=attr.name, writable=False)
                 type_hint = schema_hints.get(attr.name, {})
                 # type hints take precedence
-                if 'type' in type_hint:
-                    schema['type'] = type_hint['type']
-                    if schema['type'] == 'categorical' and 'categories' in type_hint:
-                        schema['categories'] = type_hint['categories']
+                if "type" in type_hint:
+                    schema["type"] = type_hint["type"]
+                    if schema["type"] == "categorical" and "categories" in type_hint:
+                        schema["categories"] = type_hint["categories"]
                 else:
                     schema.update(dtype_to_schema(attr.dtype))
                 cols.append(schema)
 
             annotations[ax] = dict(columns=cols)
 
-            if 'index' in schema_hints:
-                annotations[ax].update({'index': schema_hints['index']})
+            if "index" in schema_hints:
+                annotations[ax].update({"index": schema_hints["index"]})
 
         obs_layout = []
         embeddings = self.get_embedding_names()
         for ename in embeddings:
             A = self.open_array(f"emb/{ename}")
-            obs_layout.append({
-                'name': ename,
-                'type': A.dtype.name,
-                'dims': [f'{ename}_{d}' for d in range(0, A.ndim)]
-            })
+            obs_layout.append({"name": ename, "type": A.dtype.name, "dims": [f"{ename}_{d}" for d in range(0, A.ndim)]})
 
-        schema = {
-            'dataframe': dataframe,
-            'annotations': annotations,
-            'layout': {'obs': obs_layout}
-        }
+        schema = {"dataframe": dataframe, "annotations": annotations, "layout": {"obs": obs_layout}}
         return schema
 
     def annotation_to_fbs_matrix(self, axis, fields=None, labels=None):
-        with ServerTiming.time(f'annotations.{axis}.query'):
+        with ServerTiming.time(f"annotations.{axis}.query"):
             A = self.open_array(str(axis))
             if fields is not None and len(fields) > 0:
                 try:
@@ -326,7 +306,7 @@ class CxgAdaptor(DataAdaptor):
                     obs_names = self.get_obs_names()
                     df = df.join(labels, obs_names)
 
-        with ServerTiming.time(f'annotations.{axis}.encode'):
+        with ServerTiming.time(f"annotations.{axis}.encode"):
             fbs = encode_matrix_fbs(df, col_idx=df.columns)
 
         return fbs
@@ -337,7 +317,7 @@ class CxgAdaptor(DataAdaptor):
         if boolarray is None:
             return slice(None)
         assert type(boolarray) == np.ndarray
-        assert(boolarray.dtype) == bool
+        assert (boolarray.dtype) == bool
 
         selector = np.nonzero(boolarray)[0]
 

server/test/__init__.py

@@ -0,0 +1,50 @@
+import shutil
+import tempfile
+from os import path
+
+import pandas as pd
+
+from server.common.annotations import AnnotationsLocalFile
+from server.common.data_locator import DataLocator
+from server.data_common.fbs.matrix import encode_matrix_fbs
+from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataType
+
+
+def data_with_tmp_annotations(ext: MatrixDataType, annotations_fixture=False):
+    tmp_dir = tempfile.mkdtemp()
+    annotations_file = path.join(tmp_dir, "test_annotations.csv")
+    if annotations_fixture:
+        shutil.copyfile(f"test/test_datasets/pbmc3k-annotations.csv", annotations_file)
+    args = {
+        "layout": ["umap"],
+        "max_category_items": 100,
+        "obs_names": None,
+        "var_names": None,
+        "diffexp_lfc_cutoff": 0.01,
+    }
+    fname = {
+        MatrixDataType.H5AD: "../example-dataset/pbmc3k.h5ad",
+        MatrixDataType.CXG: "test/test_datasets/pbmc3k.cxg",
+    }[ext]
+    data_locator = DataLocator(fname)
+    data = MatrixDataLoader(data_locator.abspath()).open(args)
+    annotations = AnnotationsLocalFile(None, annotations_file)
+    return data, tmp_dir, annotations
+
+
+def make_fbs(data):
+    df = pd.DataFrame(data)
+    return encode_matrix_fbs(matrix=df, row_idx=None, col_idx=df.columns)
+
+
+def skip_if(condition, reason: str):
+    def decorator(f):
+        def wraps(self, *args, **kwargs):
+            if condition(self):
+                self.skipTest(reason)
+            else:
+                f(self, *args, **kwargs)
+
+        return wraps
+
+    return decorator

server/test/test_anndata_adaptor.py

@@ -3,7 +3,7 @@ from os import path
 import pytest
 import time
 import unittest
-import decode_fbs
+import server.test.decode_fbs as decode_fbs
 from parameterized import parameterized_class
 
 import numpy as np

server/test/test_api.py

@@ -1,17 +1,24 @@
+import shutil
+import time
+import unittest
 from http import HTTPStatus
 from subprocess import Popen
-import unittest
-import time
 
+import pandas as pd
 import requests
 
-import decode_fbs
-
+import server.test.decode_fbs as decode_fbs
+from server.test import skip_if, data_with_tmp_annotations, make_fbs
+from server.data_common.matrix_loader import MatrixDataType
 
 BAD_FILTER = {"filter": {"obs": {"annotation_value": [{"name": "xyz"}]}}}
 
+# TODO (mweiden): remove ANNOTATIONS_ENABLED and Annotation subclasses when annotations are no longer experimental
+# TODO (mweiden): remove MATRIX_DATA_TYPE and skip_if when user annotations for the CXG format is complete
+
 
 class EndPoints(object):
+    ANNOTATIONS_ENABLED = False
 
     def setUp(self):
         self.session = requests.Session()
@@ -25,7 +32,9 @@ class EndPoints(object):
         result_data = result.json()
         self.assertEqual(result_data["schema"]["dataframe"]["nObs"], 2638)
         self.assertEqual(len(result_data["schema"]["annotations"]["obs"]), 2)
-        self.assertEqual(len(result_data["schema"]["annotations"]["obs"]["columns"]), 5)
+        self.assertEqual(
+            len(result_data["schema"]["annotations"]["obs"]["columns"]), 6 if self.ANNOTATIONS_ENABLED else 5
+        )
 
     def test_config(self):
         endpoint = "config"
@@ -51,7 +60,7 @@ class EndPoints(object):
         self.assertIsNotNone(df["columns"])
         self.assertSetEqual(
             set(df["col_idx"]),
-            set(["pca_0", "pca_1", "tsne_0", "tsne_1", "umap_0", "umap_1", "draw_graph_fr_0", "draw_graph_fr_1"]),
+            {"pca_0", "pca_1", "tsne_0", "tsne_1", "umap_0", "umap_1", "draw_graph_fr_0", "draw_graph_fr_1"},
         )
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
@@ -71,14 +80,21 @@ class EndPoints(object):
         self.assertEqual(result.headers["Content-Type"], "application/octet-stream")
         df = decode_fbs.decode_matrix_FBS(result.content)
         self.assertEqual(df["n_rows"], 2638)
-        self.assertEqual(df["n_cols"], 5)
+        self.assertEqual(df["n_cols"], 6 if self.ANNOTATIONS_ENABLED else 5)
         self.assertIsNotNone(df["columns"])
-        self.assertIsNotNone(df["col_idx"])
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
         obs_index_col_name = self.schema["schema"]["annotations"]["obs"]["index"]
-        self.assertListEqual(df["col_idx"], [obs_index_col_name, "n_genes", "percent_mito", "n_counts", "louvain"])
+        self.assertListEqual(
+            df["col_idx"],
+            [obs_index_col_name, "n_genes", "percent_mito", "n_counts", "louvain"]
+            + (["cluster-test"] if self.ANNOTATIONS_ENABLED else []),
+        )
 
+    @skip_if(
+        lambda slf: hasattr(slf, "MATRIX_DATA_TYPE") and slf.MATRIX_DATA_TYPE == MatrixDataType.CXG,
+        "CXG file annotations are not feature-complete!",
+    )
     def test_get_annotations_obs_keys_fbs(self):
         endpoint = "annotations/obs"
         query = "annotation-name=n_genes&annotation-name=percent_mito"
@@ -91,7 +107,6 @@ class EndPoints(object):
         self.assertEqual(df["n_rows"], 2638)
         self.assertEqual(df["n_cols"], 2)
         self.assertIsNotNone(df["columns"])
-        self.assertIsNotNone(df["col_idx"])
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
         self.assertListEqual(df["col_idx"], ["n_genes", "percent_mito"])
@@ -144,7 +159,6 @@ class EndPoints(object):
         self.assertEqual(df["n_rows"], 1838)
         self.assertEqual(df["n_cols"], 2)
         self.assertIsNotNone(df["columns"])
-        self.assertIsNotNone(df["col_idx"])
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
         var_index_col_name = self.schema["schema"]["annotations"]["var"]["index"]
@@ -162,7 +176,6 @@ class EndPoints(object):
         self.assertEqual(df["n_rows"], 1838)
         self.assertEqual(df["n_cols"], 1)
         self.assertIsNotNone(df["columns"])
-        self.assertIsNotNone(df["col_idx"])
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
         self.assertListEqual(df["col_idx"], ["n_cells"])
@@ -215,7 +228,6 @@ class EndPoints(object):
         self.assertEqual(df["n_rows"], 2638)
         self.assertEqual(df["n_cols"], 3)
         self.assertIsNotNone(df["columns"])
-        self.assertIsNotNone(df["col_idx"])
         self.assertIsNone(df["row_idx"])
         self.assertEqual(len(df["columns"]), df["n_cols"])
         self.assertListEqual(df["col_idx"].tolist(), [0, 1, 4])
@@ -240,6 +252,77 @@ class EndPoints(object):
         result = self.session.get(url)
         self.assertEqual(result.status_code, HTTPStatus.OK)
 
+    @staticmethod
+    def _setUpClass(child_class, start_command):
+        child_class.ps = Popen(start_command)
+        child_class.session = requests.Session()
+        for i in range(90):
+            try:
+                result = child_class.session.get(f"{child_class.URL_BASE}schema")
+                child_class.schema = result.json()
+            except requests.exceptions.ConnectionError:
+                time.sleep(1)
+
+    @staticmethod
+    def _tearDownClass(child_class):
+        try:
+            child_class.ps.terminate()
+        except ProcessLookupError:
+            pass
+
+
+class EndPointsAnnotations(EndPoints):
+    def test_get_schema_existing_writable(self):
+        self._test_get_schema_writable("cluster-test")
+
+    @skip_if(lambda slf: slf.MATRIX_DATA_TYPE == MatrixDataType.CXG, "CXG file annotations are not feature-complete!")
+    def test_get_user_annotations_existing_obs_keys_fbs(self):
+        self._test_get_user_annotations_obs_keys_fbs(
+            "cluster-test", {"unassigned", "one", "two", "three", "four", "five"},
+        )
+
+    @skip_if(lambda slf: slf.MATRIX_DATA_TYPE == MatrixDataType.CXG, "CXG file annotations are not feature-complete!")
+    def test_put_user_annotations_obs_fbs(self):
+        endpoint = "annotations/obs"
+        query = "annotation-collection-name=test_annotations"
+        url = f"{self.URL_BASE}{endpoint}?{query}"
+        n_rows = self.data.get_shape()[0]
+        fbs = make_fbs({"cat_A": pd.Series(["label_A" for l in range(0, n_rows)], dtype="category")})
+        result = self.session.put(url, data=fbs)
+        self.assertEqual(result.status_code, HTTPStatus.OK)
+        self.assertEqual(result.headers["Content-Type"], "application/json")
+        self.assertEqual(result.json(), {"status": "OK"})
+        self._test_get_schema_writable("cat_A")
+        self._test_get_user_annotations_obs_keys_fbs("cat_A", {"label_A"})
+
+    def _test_get_user_annotations_obs_keys_fbs(self, annotation_name, columns):
+        endpoint = "annotations/obs"
+        query = f"annotation-name={annotation_name}"
+        url = f"{self.URL_BASE}{endpoint}?{query}"
+        header = {"Accept": "application/octet-stream"}
+        result = self.session.get(url, headers=header)
+        self.assertEqual(result.status_code, HTTPStatus.OK)
+        self.assertEqual(result.headers["Content-Type"], "application/octet-stream")
+        df = decode_fbs.decode_matrix_FBS(result.content)
+        self.assertEqual(df["n_rows"], 2638)
+        self.assertEqual(df["n_cols"], 1)
+        self.assertListEqual(df["col_idx"], [annotation_name])
+        self.assertEqual(set(df["columns"][0]), columns)
+        self.assertIsNone(df["row_idx"])
+        self.assertEqual(len(df["columns"]), df["n_cols"])
+
+    def _test_get_schema_writable(self, cluster_name):
+        endpoint = "schema"
+        url = f"{self.URL_BASE}{endpoint}"
+        result = self.session.get(url)
+        self.assertEqual(result.status_code, HTTPStatus.OK)
+        self.assertEqual(result.headers["Content-Type"], "application/json")
+        result_data = result.json()
+        columns = result_data["schema"]["annotations"]["obs"]["columns"]
+        matching_columns = [c for c in columns if c["name"] == cluster_name]
+        self.assertEqual(len(matching_columns), 1)
+        self.assertTrue(matching_columns[0]["writable"])
+
 
 class EndPointsAnndata(unittest.TestCase, EndPoints):
     """Test Case for endpoints"""
@@ -251,7 +334,8 @@ class EndPointsAnndata(unittest.TestCase, EndPoints):
 
     @classmethod
     def setUpClass(cls):
-        cls.ps = Popen(
+        cls._setUpClass(
+            cls,
             [
                 "cellxgene",
                 "--no-upgrade-check",
@@ -260,22 +344,16 @@ class EndPointsAnndata(unittest.TestCase, EndPoints):
                 "--verbose",
                 "--port",
                 str(cls.PORT),
-            ]
+            ],
         )
-        cls.session = requests.Session()
-        for i in range(90):
-            try:
-                result = cls.session.get(f"{cls.URL_BASE}schema")
-                cls.schema = result.json()
-            except requests.exceptions.ConnectionError:
-                time.sleep(1)
 
     @classmethod
     def tearDownClass(cls):
-        try:
-            cls.ps.terminate()
-        except ProcessLookupError:
-            pass
+        cls._tearDownClass(cls)
+
+    @property
+    def annotations_enabled(self):
+        return False
 
 
 class EndPointsCxg(unittest.TestCase, EndPoints):
@@ -288,28 +366,91 @@ class EndPointsCxg(unittest.TestCase, EndPoints):
 
     @classmethod
     def setUpClass(cls):
-        cls.ps = Popen(
+        cls._setUpClass(
+            cls,
             [
                 "cellxgene",
                 "--no-upgrade-check",
                 "launch",
-                "../example-dataset/pbmc3k.cxg",
+                "test/test_datasets/pbmc3k.cxg",
                 "--verbose",
                 "--port",
                 str(cls.PORT),
-            ]
+            ],
         )
-        cls.session = requests.Session()
-        for i in range(90):
-            try:
-                result = cls.session.get(f"{cls.URL_BASE}schema")
-                cls.schema = result.json()
-            except requests.exceptions.ConnectionError:
-                time.sleep(1)
 
     @classmethod
     def tearDownClass(cls):
-        try:
-            cls.ps.terminate()
-        except ProcessLookupError:
-            pass
+        cls._tearDownClass(cls)
+
+
+class EndPointsAnndataAnnotations(unittest.TestCase, EndPointsAnnotations):
+    """Test Case for endpoints"""
+
+    PORT = 5012
+    LOCAL_URL = f"http://127.0.0.1:{PORT}/"
+    VERSION = "v0.2"
+    URL_BASE = f"{LOCAL_URL}api/{VERSION}/"
+    ANNOTATIONS_ENABLED = True
+    MATRIX_DATA_TYPE = MatrixDataType.H5AD
+
+    @classmethod
+    def setUpClass(cls):
+        cls.data, cls.tmp_dir, cls.annotations = data_with_tmp_annotations(
+            MatrixDataType.H5AD, annotations_fixture=True
+        )
+        cls._setUpClass(
+            cls,
+            [
+                "cellxgene",
+                "--no-upgrade-check",
+                "launch",
+                "--experimental-annotations",
+                "--experimental-annotations-file",
+                cls.annotations.output_file,
+                "--verbose",
+                "--port",
+                str(cls.PORT),
+                cls.data.get_location(),
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        shutil.rmtree(cls.tmp_dir)
+        cls._tearDownClass(cls)
+
+
+class EndPointsCxgAnnotations(unittest.TestCase, EndPointsAnnotations):
+    """Test Case for endpoints"""
+
+    PORT = 5013
+    LOCAL_URL = f"http://127.0.0.1:{PORT}/"
+    VERSION = "v0.2"
+    URL_BASE = f"{LOCAL_URL}api/{VERSION}/"
+    ANNOTATIONS_ENABLED = True
+    MATRIX_DATA_TYPE = MatrixDataType.CXG
+
+    @classmethod
+    def setUpClass(cls):
+        cls.data, cls.tmp_dir, cls.annotations = data_with_tmp_annotations(MatrixDataType.CXG, annotations_fixture=True)
+        cls._setUpClass(
+            cls,
+            [
+                "cellxgene",
+                "--no-upgrade-check",
+                "launch",
+                "--experimental-annotations",
+                "--experimental-annotations-file",
+                cls.annotations.output_file,
+                "--verbose",
+                "--port",
+                str(cls.PORT),
+                cls.data.get_location(),
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        shutil.rmtree(cls.tmp_dir)
+        cls._tearDownClass(cls)

server/test/test_fbs.py

@@ -3,7 +3,7 @@ import pandas as pd
 import numpy as np
 from scipy import sparse
 
-import decode_fbs
+import server.test.decode_fbs as decode_fbs
 from server.data_common.fbs.matrix import encode_matrix_fbs, decode_matrix_fbs
 
 

server/test/test_nan_anndata_adaptor.py

@@ -3,7 +3,7 @@ import unittest
 import warnings
 import math
 
-import decode_fbs
+import server.test.decode_fbs as decode_fbs
 
 from server.data_anndata.anndata_adaptor import AnndataAdaptor
 from server.common.errors import FilterError

server/test/test_nan_rest.py

@@ -4,7 +4,7 @@ import unittest
 import time
 import math
 
-import decode_fbs
+import server.test.decode_fbs as decode_fbs
 
 import requests
 

server/test/test_writable_annotation.py

@@ -1,42 +1,23 @@
 import json
 from os import path, listdir
 import unittest
-import decode_fbs
-import tempfile
+import server.test.decode_fbs as decode_fbs
 import shutil
 
 import numpy as np
 import pandas as pd
 
-from server.data_anndata.anndata_adaptor import AnndataAdaptor
-from server.data_common.fbs.matrix import encode_matrix_fbs
-from server.common.data_locator import DataLocator
-from server.common.annotations import AnnotationsLocalFile
 from server.common.rest import schema_get_helper, annotations_put_fbs_helper
+from server.test import data_with_tmp_annotations, make_fbs
+from server.data_common.matrix_loader import MatrixDataType
 
 
 class WritableAnnotationTest(unittest.TestCase):
     def setUp(self):
-        self.tmpDir = tempfile.mkdtemp()
-        self.annotations_file = path.join(self.tmpDir, "test_annotations.csv")
-        args = {
-            "layout": ["umap"],
-            "max_category_items": 100,
-            "obs_names": None,
-            "var_names": None,
-            "diffexp_lfc_cutoff": 0.01,
-        }
-        fname = "../example-dataset/pbmc3k.h5ad"
-        data_locator = DataLocator(fname)
-        self.data = AnndataAdaptor(data_locator, args)
-        self.annotations = AnnotationsLocalFile(None, self.annotations_file)
+        self.data, self.tmp_dir, self.annotations = data_with_tmp_annotations(MatrixDataType.H5AD)
 
     def tearDown(self):
-        shutil.rmtree(self.tmpDir)
-
-    def make_fbs(self, data):
-        df = pd.DataFrame(data)
-        return encode_matrix_fbs(matrix=df, row_idx=None, col_idx=df.columns)
+        shutil.rmtree(self.tmp_dir)
 
     def annotation_put_fbs(self, fbs):
         annotations_put_fbs_helper(self.data, self.annotations, fbs)
@@ -45,8 +26,8 @@ class WritableAnnotationTest(unittest.TestCase):
 
     def test_error_checks(self):
         # verify that the expected errors are generated
-        n_rows = self.data.data.obs.shape[0]
-        fbs_bad = self.make_fbs({"louvain": pd.Series(["undefined" for l in range(0, n_rows)], dtype="category")})
+        n_rows = self.data.get_shape()[0]
+        fbs_bad = make_fbs({"louvain": pd.Series(["undefined" for l in range(0, n_rows)], dtype="category")})
 
         # ensure we catch attempt to overwrite non-writable data
         with self.assertRaises(KeyError):
@@ -54,8 +35,8 @@ class WritableAnnotationTest(unittest.TestCase):
 
     def test_write_to_file(self):
         # verify the file is written as expected
-        n_rows = self.data.data.obs.shape[0]
-        fbs = self.make_fbs(
+        n_rows = self.data.get_shape()[0]
+        fbs = make_fbs(
             {
                 "cat_A": pd.Series(["label_A" for l in range(0, n_rows)], dtype="category"),
                 "cat_B": pd.Series(["label_B" for l in range(0, n_rows)], dtype="category"),
@@ -63,8 +44,8 @@ class WritableAnnotationTest(unittest.TestCase):
         )
         res = self.annotation_put_fbs(fbs)
         self.assertEqual(res, json.dumps({"status": "OK"}))
-        self.assertTrue(path.exists(self.annotations_file))
-        df = pd.read_csv(self.annotations_file, index_col=0, header=0, comment="#")
+        self.assertTrue(path.exists(self.annotations.output_file))
+        df = pd.read_csv(self.annotations.output_file, index_col=0, header=0, comment="#")
         self.assertEqual(df.shape, (n_rows, 2))
         self.assertEqual(set(df.columns), {"cat_A", "cat_B"})
         self.assertTrue(self.data.original_obs_index.equals(df.index))
@@ -72,7 +53,7 @@ class WritableAnnotationTest(unittest.TestCase):
         self.assertTrue(np.all(df["cat_B"] == ["label_B" for l in range(0, n_rows)]))
 
         # verify complete overwrite on second attempt, AND rotation occurs
-        fbs = self.make_fbs(
+        fbs = make_fbs(
             {
                 "cat_A": pd.Series(["label_A1" for l in range(0, n_rows)], dtype="category"),
                 "cat_C": pd.Series(["label_C" for l in range(0, n_rows)], dtype="category"),
@@ -80,14 +61,14 @@ class WritableAnnotationTest(unittest.TestCase):
         )
         res = self.annotation_put_fbs(fbs)
         self.assertEqual(res, json.dumps({"status": "OK"}))
-        self.assertTrue(path.exists(self.annotations_file))
-        df = pd.read_csv(self.annotations_file, index_col=0, header=0, comment="#")
+        self.assertTrue(path.exists(self.annotations.output_file))
+        df = pd.read_csv(self.annotations.output_file, index_col=0, header=0, comment="#")
         self.assertEqual(set(df.columns), {"cat_A", "cat_C"})
         self.assertTrue(np.all(df["cat_A"] == ["label_A1" for l in range(0, n_rows)]))
         self.assertTrue(np.all(df["cat_C"] == ["label_C" for l in range(0, n_rows)]))
 
         # rotation
-        name, ext = path.splitext(self.annotations_file)
+        name, ext = path.splitext(self.annotations.output_file)
         backup_dir = f"{name}-backups"
         self.assertTrue(path.isdir(backup_dir))
         found_files = listdir(backup_dir)
@@ -95,8 +76,8 @@ class WritableAnnotationTest(unittest.TestCase):
 
     def test_file_rotation_to_max_9(self):
         # verify we stop rotation at 9
-        n_rows = self.data.data.obs.shape[0]
-        fbs = self.make_fbs(
+        n_rows = self.data.get_shape()[0]
+        fbs = make_fbs(
             {
                 "cat_A": pd.Series(["label_A" for l in range(0, n_rows)], dtype="category"),
                 "cat_B": pd.Series(["label_B" for l in range(0, n_rows)], dtype="category"),
@@ -106,7 +87,7 @@ class WritableAnnotationTest(unittest.TestCase):
             res = self.annotation_put_fbs(fbs)
             self.assertEqual(res, json.dumps({"status": "OK"}))
 
-        name, ext = path.splitext(self.annotations_file)
+        name, ext = path.splitext(self.annotations.output_file)
         backup_dir = f"{name}-backups"
         self.assertTrue(path.isdir(backup_dir))
         found_files = listdir(backup_dir)
@@ -116,8 +97,8 @@ class WritableAnnotationTest(unittest.TestCase):
         # verify that OBS PUTs (annotation_put_fbs) are accessible via
         # GET (annotation_to_fbs_matrix)
 
-        n_rows = self.data.data.obs.shape[0]
-        fbs = self.make_fbs(
+        n_rows = self.data.get_shape()[0]
+        fbs = make_fbs(
             {
                 "cat_A": pd.Series(["label_A" for l in range(0, n_rows)], dtype="category"),
                 "cat_B": pd.Series(["label_B" for l in range(0, n_rows)], dtype="category"),