add accessLevel info in datasetAlleleResponse

cgbeacon2/models/dataset_response.py

@@ -4,14 +4,30 @@
 class DatasetAlleleResponse:
     """Create a Beacon Dataset Allele Response object to be returned by Beacon Response"""
 
-    def __init__(self, dataset_id, variants):
-        self.datasetId = dataset_id
-        samples, alleles = self._sample_allele_count(dataset_id, variants)
-        self.sampleCount = samples
-        self.callCount = alleles
+    def __init__(self, dataset, variants):
+        self.datasetId = dataset["_id"]
+        n_samples, n_alleles, n_variants = self._sample_allele_variant_count(
+            self.datasetId, variants
+        )
+        self.sampleCount = n_samples
+        self.callCount = n_alleles
+        self.variantCount = n_variants
+        self.info = self._info(dataset)
         self.exists = True if self.sampleCount > 0 else False
 
-    def _sample_allele_count(self, dataset_id, variants):
+    def _info(self, dataset):
+        """Provides additional info regarding a dataset object
+
+        Accepts:
+            dataset(dict)
+
+        Returns:
+            info(dict)
+        """
+        info = dict(accessType=dataset["authlevel"].upper())
+        return info
+
+    def _sample_allele_variant_count(self, dataset_id, variants):
         """Counts samples and allelic calls for one or more variants
 
         Accepts:
@@ -23,11 +39,13 @@ def _sample_allele_count(self, dataset_id, variants):
         """
         n_samples = 0
         n_calls = 0
+        n_variants = 0
         for variant_obj in variants:
             if dataset_id in variant_obj.get("datasetIds"):
                 if variant_obj["datasetIds"][dataset_id].get("samples"):
                     n_samples += len(
                         variant_obj["datasetIds"][dataset_id]["samples"].keys()
                     )
+                n_variants += 1
                 n_calls += variant_obj["call_count"]
-        return n_samples, n_calls
+        return n_samples, n_calls, n_variants

cgbeacon2/server/blueprints/api_v1/controllers.py

@@ -308,7 +308,7 @@ def create_ds_allele_response(response_type, req_dsets, variants):
     exists = False
 
     all_dsets = current_app.db["dataset"].find()
-    all_dsets = [ds["_id"] for ds in all_dsets]
+    all_dsets = {ds["_id"]: ds for ds in all_dsets}
 
     if len(req_dsets) == 0:  # if query didn't specify any dataset
         # Use all datasets present in this beacon
@@ -319,7 +319,8 @@ def create_ds_allele_response(response_type, req_dsets, variants):
         if not ds in all_dsets:
             LOG.info(f"Provided dataset {ds} could not be found in database")
             continue
-        ds_response = DatasetAlleleResponse(ds, variants).__dict__
+
+        ds_response = DatasetAlleleResponse(all_dsets[ds], variants).__dict__
 
         # collect responses according to the type of response requested
         if (

tests/server/blueprints/api_v1/test_views_query_no_auth.py

@@ -90,6 +90,13 @@ def test_get_request_exact_position_snv_return_ALL(
         if ds_level_result["exists"] is True:
             assert ds_level_result["callCount"] > 0
 
+        # Variant count should also be a positive number if there is positive match
+        if ds_level_result["exists"] is True:
+            assert ds_level_result["variantCount"] > 0
+
+        # Dataset info should be available and containing the expected info
+        assert ds_level_result["info"] == {"accessType": "PUBLIC"}
+
 
 def test_get_request_exact_position_snv_return_HIT(
     mock_app, test_snv, public_dataset, public_dataset_no_variants