Repository: CartoDB/cartoframes
Pull request #1234 (2019-11-22 17:03:05)

Validate BQ operations
Solves https://github.com/CartoDB/cartoframes/issues/1228

Tasks:
- Add new props in catalog entities
- Validate if dataset/geography is ready before enrichment and download


Patch:
diff --git a/cartoframes/data/observatory/catalog/entity.py b/cartoframes/data/observatory/catalog/entity.py
index 71b7561c..210a8111 100644
--- a/cartoframes/data/observatory/catalog/entity.py
+++ b/cartoframes/data/observatory/catalog/entity.py
@@ -15,13 +15,14 @@ except ImportError:
     ABC = ABCMeta('ABC', (object,), {'__slots__': ()})
 
 _WORKING_PROJECT = 'carto-do-customers'
+_PLATFORM_BQ = 'bq'
 
 
 class CatalogEntity(ABC):
 
     id_field = 'id'
     entity_repo = None
-    export_excluded_fields = ['summary_json']
+    export_excluded_fields = ['summary_json', 'available_in']
 
     def __init__(self, data):
         self.data = data
@@ -74,6 +75,9 @@ class CatalogEntity(ABC):
         return self.id
 
     def _download(self, credentials=None):
+        if not self._is_available_in('bq'):
+            raise CartoException('{} is not ready for Download. Please, contact us for more information.'.format(self))
+
         credentials = self._get_credentials(credentials)
         user_dataset = credentials.get_do_user_dataset()
         bq_client = _get_bigquery_client(_WORKING_PROJECT, credentials)
@@ -91,6 +95,9 @@ class CatalogEntity(ABC):
 
         return file_path
 
+    def _is_available_in(self, platform=_PLATFORM_BQ):
+        return self.data['available_in'] and platform in self.data['available_in']
+
     def _get_credentials(self, credentials=None):
         _credentials = credentials or defaults.get_default_credentials()
 
diff --git a/cartoframes/data/observatory/catalog/geography.py b/cartoframes/data/observatory/catalog/geography.py
index af2e303d..c45978d2 100644
--- a/cartoframes/data/observatory/catalog/geography.py
+++ b/cartoframes/data/observatory/catalog/geography.py
@@ -65,6 +65,12 @@ class Geography(CatalogEntity):
 
         return self.data['geom_coverage']
 
+    @property
+    def geom_type(self):
+        """Info about the type of geometry of this geography."""
+
+        return self.data['geom_type']
+
     @property
     def update_frequency(self):
         """Frequency in which the geography is updated."""
diff --git a/cartoframes/data/observatory/catalog/repository/dataset_repo.py b/cartoframes/data/observatory/catalog/repository/dataset_repo.py
index cac4420e..32414a89 100644
--- a/cartoframes/data/observatory/catalog/repository/dataset_repo.py
+++ b/cartoframes/data/observatory/catalog/repository/dataset_repo.py
@@ -55,7 +55,8 @@ class DatasetRepository(EntityRepository):
             'update_frequency': self._normalize_field(row, 'update_frequency'),
             'version': self._normalize_field(row, 'version'),
             'is_public_data': self._normalize_field(row, 'is_public_data'),
-            'summary_json': self._normalize_field(row, 'summary_json')
+            'summary_json': self._normalize_field(row, 'summary_json'),
+            'available_in': self._normalize_field(row, 'available_in')
         }
 
     def get_datasets_for_geographies(self, geographies):
diff --git a/cartoframes/data/observatory/catalog/repository/geography_repo.py b/cartoframes/data/observatory/catalog/repository/geography_repo.py
index e8548a65..c2438307 100644
--- a/cartoframes/data/observatory/catalog/repository/geography_repo.py
+++ b/cartoframes/data/observatory/catalog/repository/geography_repo.py
@@ -52,10 +52,12 @@ class GeographyRepository(EntityRepository):
             'provider_name': self._normalize_field(row, 'provider_name'),
             'lang': self._normalize_field(row, 'lang'),
             'geom_coverage': self._normalize_field(row, 'geom_coverage'),
+            'geom_type': self._normalize_field(row, 'geom_type'),
             'update_frequency': self._normalize_field(row, 'update_frequency'),
             'version': self._normalize_field(row, 'version'),
             'is_public_data': self._normalize_field(row, 'is_public_data'),
-            'summary_json': self._normalize_field(row, 'summary_json')
+            'summary_json': self._normalize_field(row, 'summary_json'),
+            'available_in': self._normalize_field(row, 'available_in')
         }
 
     def get_geographies_gdf(self):
diff --git a/cartoframes/data/observatory/enrichment/enrichment_service.py b/cartoframes/data/observatory/enrichment/enrichment_service.py
index e8b5cab1..b624b310 100644
--- a/cartoframes/data/observatory/enrichment/enrichment_service.py
+++ b/cartoframes/data/observatory/enrichment/enrichment_service.py
@@ -4,6 +4,7 @@ from collections import defaultdict
 
 from ..catalog.variable import Variable
 from ..catalog.dataset import Dataset
+from ..catalog.geography import Geography
 from ...clients import bigquery_client
 from ....auth import get_default_credentials
 from ....exceptions import EnrichmentException
@@ -190,9 +191,22 @@ def _prepare_variable(variable):
             Variable `id` property or Variable `slug` property
         """)
 
+    _is_available_in_bq(variable)
+
     return variable
 
 
+def _is_available_in_bq(variable):
+    dataset = Dataset.get(variable.dataset)
+    geography = Geography.get(dataset.geography)
+
+    if not (dataset._is_available_in('bq') and geography._is_available_in('bq')):
+        raise EnrichmentException("""
+            The Dataset or the Geography of the Variable '{}' is not ready for Enrichment.
+            Please, contact us for more information.
+        """.format(variable.slug))
+
+
 def get_variable_aggregations(variables, aggregation):
     return [VariableAggregation(variable, __get_aggregation(variable, aggregation)) for variable in variables]
 
diff --git a/examples/_debug/do_catalog_discovery.ipynb b/examples/_debug/do_catalog_discovery.ipynb
index 55837656..2d9579a7 100644
--- a/examples/_debug/do_catalog_discovery.ipynb
+++ b/examples/_debug/do_catalog_discovery.ipynb
@@ -203,9 +203,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cartoframes.data.observatory.dataset import CatalogDataset\n",
+    "from cartoframes.data.observatory.dataset import Dataset\n",
     "\n",
-    "isinstance(demographics_datasets[0], CatalogDataset)"
+    "isinstance(demographics_datasets[0], Dataset)"
    ]
   },
   {
@@ -276,7 +276,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "CatalogDataset.get('od_acsquantile_928a2a23').to_dict()"
+    "Dataset.get('od_acsquantile_928a2a23').to_dict()"
    ]
   },
   {
@@ -292,7 +292,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "CatalogDataset.get_list(['od_acsquantile_ae4e7c82', 'od_acs_13345497'])"
+    "Dataset.get_list(['od_acsquantile_ae4e7c82', 'od_acs_13345497'])"
    ]
   },
   {
diff --git a/examples/_debug/do_subscriptions.ipynb b/examples/_debug/do_subscriptions.ipynb
index db185e2b..62ba6bff 100644
--- a/examples/_debug/do_subscriptions.ipynb
+++ b/examples/_debug/do_subscriptions.ipynb
@@ -7,7 +7,7 @@
    "outputs": [],
    "source": [
     "from cartoframes.auth import Credentials, set_default_credentials\n",
-    "from cartoframes.data.observatory import Catalog, CatalogDataset, Geography"
+    "from cartoframes.data.observatory import Catalog, Dataset, Geography"
    ]
   },
   {
@@ -35,7 +35,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dataset = CatalogDataset.get('ags_climateandw_cffac915')"
+    "dataset = Dataset.get('ags_climateandw_cffac915')"
    ]
   },
   {
diff --git a/examples/data_enrichment/enrichment_subscription_workflow.ipynb b/examples/data_enrichment/enrichment_subscription_workflow.ipynb
index ab3797b5..207499da 100644
--- a/examples/data_enrichment/enrichment_subscription_workflow.ipynb
+++ b/examples/data_enrichment/enrichment_subscription_workflow.ipynb
@@ -29,7 +29,7 @@
     {
      "data": {
       "text/plain": [
-       "(<CatalogDataset('ags_retailpoten_ddf56a1a')>,\n",
+       "(<Dataset('ags_retailpoten_ddf56a1a')>,\n",
        " <Geography('ags_blockgroup_1c63771c')>)"
       ]
      },
@@ -39,9 +39,9 @@
     }
    ],
    "source": [
-    "from cartoframes.data.observatory import Catalog, CatalogDataset, Geography\n",
+    "from cartoframes.data.observatory import Catalog, Dataset, Geography\n",
     "\n",
-    "dataset = CatalogDataset.get(dataset_name)\n",
+    "dataset = Dataset.get(dataset_name)\n",
     "geography = Geography.get(dataset.geography)\n",
     "\n",
     "dataset, geography"
@@ -103,7 +103,7 @@
     {
      "data": {
       "text/plain": [
-       "Datasets: [<CatalogDataset('ags_retailpoten_ddf56a1a')>]\n",
+       "Datasets: [<Dataset('ags_retailpoten_ddf56a1a')>]\n",
        "Geographies: [<Geography('ags_blockgroup_1c63771c')>]"
       ]
      },
diff --git a/examples/data_enrichment/points_enrichment_private_dataset.ipynb b/examples/data_enrichment/points_enrichment_private_dataset.ipynb
index e0346217..9ec927b8 100644
--- a/examples/data_enrichment/points_enrichment_private_dataset.ipynb
+++ b/examples/data_enrichment/points_enrichment_private_dataset.ipynb
@@ -111,9 +111,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cartoframes.data.observatory import CatalogDataset\n",
+    "from cartoframes.data.observatory import Dataset\n",
     "\n",
-    "dataset = CatalogDataset.get('carto-do.ags.demographics_retailpotential_usa_blockgroup_2015_yearly_2018')\n",
+    "dataset = Dataset.get('carto-do.ags.demographics_retailpotential_usa_blockgroup_2015_yearly_2018')\n",
     "variables = dataset.variables"
    ]
   },
diff --git a/examples/data_enrichment/points_enrichment_public_dataset.ipynb b/examples/data_enrichment/points_enrichment_public_dataset.ipynb
index 5ca745f3..d57084a3 100644
--- a/examples/data_enrichment/points_enrichment_public_dataset.ipynb
+++ b/examples/data_enrichment/points_enrichment_public_dataset.ipynb
@@ -456,9 +456,9 @@
     }
    ],
    "source": [
-    "from cartoframes.data.observatory import CatalogDataset, Catalog\n",
+    "from cartoframes.data.observatory import Dataset, Catalog\n",
     "\n",
-    "dataset = CatalogDataset.get('carto-do-public-data.usa_acs.demographics_acs_usa_censustractclipped_2015_5yrs_20132017')\n",
+    "dataset = Dataset.get('carto-do-public-data.usa_acs.demographics_acs_usa_censustractclipped_2015_5yrs_20132017')\n",
     "variables = dataset.variables\n",
     "variables"
    ]
diff --git a/examples/data_enrichment/polygons_enrichment_private_dataset.ipynb b/examples/data_enrichment/polygons_enrichment_private_dataset.ipynb
index 5c49aace..6233948f 100644
--- a/examples/data_enrichment/polygons_enrichment_private_dataset.ipynb
+++ b/examples/data_enrichment/polygons_enrichment_private_dataset.ipynb
@@ -99,9 +99,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cartoframes.data.observatory import CatalogDataset\n",
+    "from cartoframes.data.observatory import Dataset\n",
     "\n",
-    "dataset = CatalogDataset.get('carto-do.ags.demographics_retailpotential_usa_blockgroup_2015_yearly_2018')\n",
+    "dataset = Dataset.get('carto-do.ags.demographics_retailpotential_usa_blockgroup_2015_yearly_2018')\n",
     "variables = dataset.variables"
    ]
   },
diff --git a/examples/data_enrichment/polygons_enrichment_public_dataset.ipynb b/examples/data_enrichment/polygons_enrichment_public_dataset.ipynb
index 00451de4..6d849100 100644
--- a/examples/data_enrichment/polygons_enrichment_public_dataset.ipynb
+++ b/examples/data_enrichment/polygons_enrichment_public_dataset.ipynb
@@ -99,9 +99,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cartoframes.data.observatory import CatalogDataset\n",
+    "from cartoframes.data.observatory import Dataset\n",
     "\n",
-    "dataset = CatalogDataset.get('carto-do-public-data.usa_acs.demographics_acs_usa_censustractclipped_2015_5yrs_20132017')\n",
+    "dataset = Dataset.get('carto-do-public-data.usa_acs.demographics_acs_usa_censustractclipped_2015_5yrs_20132017')\n",
     "variables = dataset.variables"
    ]
   },


Test patch:
diff --git a/tests/unit/data/observatory/catalog/examples.py b/tests/unit/data/observatory/catalog/examples.py
index 07931cb2..6605ead7 100644
--- a/tests/unit/data/observatory/catalog/examples.py
+++ b/tests/unit/data/observatory/catalog/examples.py
@@ -35,10 +35,12 @@ db_geography1 = {
     'country_id': 'esp',
     'lang': 'esp',
     'geom_coverage': '',
+    'geom_type': '',
     'update_frequency': 'monthly',
     'version': '20190203',
     'is_public_data': True,
-    'summary_json': {}
+    'summary_json': {},
+    'available_in': ['bq']
 }
 db_geography2 = {
     'id': 'carto-do-public.tiger.geography_esp_municipalities_2019',
@@ -50,10 +52,12 @@ db_geography2 = {
     'country_id': 'esp',
     'lang': 'esp',
     'geom_coverage': '',
+    'geom_type': '',
     'update_frequency': 'monthly',
     'version': '20190203',
     'is_public_data': False,
-    'summary_json': {}
+    'summary_json': {},
+    'available_in': []
 }
 test_geography1 = Geography(db_geography1)
 test_geography2 = Geography(db_geography2)
@@ -79,7 +83,8 @@ db_dataset1 = {
     'update_frequency': 'monthly',
     'version': '20190203',
     'is_public_data': True,
-    'summary_json': {}
+    'summary_json': {},
+    'available_in': ['bq']
 }
 db_dataset2 = {
     'id': 'carto-do-public.project.basicstats-municipalities',
@@ -101,7 +106,8 @@ db_dataset2 = {
     'update_frequency': 'monthly',
     'version': '20190203',
     'is_public_data': False,
-    'summary_json': {}
+    'summary_json': {},
+    'available_in': []
 }
 test_dataset1 = Dataset(db_dataset1)
 test_dataset2 = Dataset(db_dataset2)
diff --git a/tests/unit/data/observatory/catalog/repository/test_dataset_repo.py b/tests/unit/data/observatory/catalog/repository/test_dataset_repo.py
index c91a99ce..5ad435ec 100644
--- a/tests/unit/data/observatory/catalog/repository/test_dataset_repo.py
+++ b/tests/unit/data/observatory/catalog/repository/test_dataset_repo.py
@@ -196,7 +196,8 @@ class TestDatasetRepo(object):
             'update_frequency': None,
             'version': None,
             'is_public_data': None,
-            'summary_json': None
+            'summary_json': None,
+            'available_in': None
         })])
 
         # When
diff --git a/tests/unit/data/observatory/catalog/repository/test_geography_repo.py b/tests/unit/data/observatory/catalog/repository/test_geography_repo.py
index 5c6172da..f1317220 100644
--- a/tests/unit/data/observatory/catalog/repository/test_geography_repo.py
+++ b/tests/unit/data/observatory/catalog/repository/test_geography_repo.py
@@ -198,10 +198,12 @@ class TestGeographyRepo(object):
             'country_id': None,
             'lang': None,
             'geom_coverage': None,
+            'geom_type': None,
             'update_frequency': None,
             'version': None,
             'is_public_data': None,
-            'summary_json': None
+            'summary_json': None,
+            'available_in': None
         })])
 
         # When
diff --git a/tests/unit/data/observatory/catalog/test_dataset.py b/tests/unit/data/observatory/catalog/test_dataset.py
index cf0cd5d8..2c16a330 100644
--- a/tests/unit/data/observatory/catalog/test_dataset.py
+++ b/tests/unit/data/observatory/catalog/test_dataset.py
@@ -143,7 +143,8 @@ class TestDataset(object):
     def test_dataset_is_exported_as_dict(self):
         # Given
         dataset = Dataset(db_dataset1)
-        expected_dict = {key: value for key, value in db_dataset1.items() if key != 'summary_json'}
+        excluded_fields = ['summary_json', 'available_in']
+        expected_dict = {key: value for key, value in db_dataset1.items() if key not in excluded_fields}
 
         # When
         dataset_dict = dataset.to_dict()
@@ -266,6 +267,28 @@ class TestDataset(object):
 
         assert response == file_path
 
+    @patch.object(DatasetRepository, 'get_by_id')
+    @patch('cartoframes.data.observatory.catalog.entity._get_bigquery_client')
+    def test_dataset_not_available_in_bq_download_fails(self, mocked_bq_client, mocked_repo):
+        # mock dataset
+        mocked_repo.return_value = test_dataset2
+
+        # mock big query client
+        file_path = 'fake_path'
+        mocked_bq_client.return_value = BigQueryClientMock(file_path)
+
+        # test
+        username = 'fake_user'
+        credentials = Credentials(username, '1234')
+
+        dataset = Dataset.get(test_dataset2.id)
+
+        with pytest.raises(CartoException) as e:
+            dataset.download(credentials)
+
+        error = '{} is not ready for Download. Please, contact us for more information.'.format(dataset)
+        assert str(e.value) == error
+
     @patch.object(DatasetRepository, 'get_by_id')
     @patch('cartoframes.data.observatory.catalog.entity._get_bigquery_client')
     def test_dataset_download_raises_with_nonpurchased(self, mocked_bq_client, mocked_repo):
@@ -399,3 +422,16 @@ class TestDataset(object):
 
         # Then
         assert str(e.value) == '`credentials` must be a Credentials class instance'
+
+    def test_dataset_is_available_in(self):
+        dataset_in_bq = Dataset(db_dataset1)
+        dataset_not_in_bq = Dataset(db_dataset2)
+
+        assert dataset_in_bq._is_available_in('bq')
+        assert not dataset_not_in_bq._is_available_in('bq')
+
+    def test_dataset_is_available_in_with_empty_field(self):
+        db_dataset = dict(db_dataset1)
+        db_dataset['available_in'] = None
+        dataset_null = Dataset(db_dataset)
+        assert not dataset_null._is_available_in('bq')
diff --git a/tests/unit/data/observatory/catalog/test_geography.py b/tests/unit/data/observatory/catalog/test_geography.py
index 93f8dd43..cb952215 100644
--- a/tests/unit/data/observatory/catalog/test_geography.py
+++ b/tests/unit/data/observatory/catalog/test_geography.py
@@ -120,7 +120,8 @@ class TestGeography(object):
     def test_geography_is_exported_as_dict(self):
         # Given
         geography = Geography(db_geography1)
-        expected_dict = {key: value for key, value in db_geography1.items() if key != 'summary_json'}
+        excluded_fields = ['summary_json', 'available_in']
+        expected_dict = {key: value for key, value in db_geography1.items() if key not in excluded_fields}
 
         # When
         geography_dict = geography.to_dict()
@@ -243,6 +244,28 @@ class TestGeography(object):
 
         assert response == file_path
 
+    @patch.object(GeographyRepository, 'get_by_id')
+    @patch('cartoframes.data.observatory.catalog.entity._get_bigquery_client')
+    def test_geography_not_available_in_bq_download_fails(self, mocked_bq_client, mocked_repo):
+        # mock geography
+        mocked_repo.return_value = test_geography2
+
+        # mock big query client
+        file_path = 'fake_path'
+        mocked_bq_client.return_value = BigQueryClientMock(file_path)
+
+        # test
+        username = 'fake_user'
+        credentials = Credentials(username, '1234')
+
+        geography = Geography.get(test_geography2.id)
+
+        with pytest.raises(CartoException) as e:
+            geography.download(credentials)
+
+        error = '{} is not ready for Download. Please, contact us for more information.'.format(geography)
+        assert str(e.value) == error
+
     @patch.object(GeographyRepository, 'get_by_id')
     @patch('cartoframes.data.observatory.catalog.entity._get_bigquery_client')
     def test_geography_download_raises_with_nonpurchased(self, mocked_bq_client, mocked_repo):
@@ -377,3 +400,16 @@ class TestGeography(object):
 
         # Then
         assert str(e.value) == '`credentials` must be a Credentials class instance'
+
+    def test_geography_is_available_in(self):
+        geography_in_bq = Geography(db_geography1)
+        geography_not_in_bq = Geography(db_geography2)
+
+        assert geography_in_bq._is_available_in('bq')
+        assert not geography_not_in_bq._is_available_in('bq')
+
+    def test_geography_is_available_in_with_empty_field(self):
+        db_geography = dict(db_geography1)
+        db_geography['available_in'] = None
+        geography_null = Geography(db_geography)
+        assert not geography_null._is_available_in('bq')
diff --git a/tests/unit/data/observatory/enrichment/test_enrichment_service.py b/tests/unit/data/observatory/enrichment/test_enrichment_service.py
index 7ea4b45d..f8969b26 100644
--- a/tests/unit/data/observatory/enrichment/test_enrichment_service.py
+++ b/tests/unit/data/observatory/enrichment/test_enrichment_service.py
@@ -1,3 +1,4 @@
+import pytest
 import pandas as pd
 from shapely.geometry.point import Point
 from shapely.geometry.polygon import Polygon
@@ -5,8 +6,10 @@ from shapely.geometry.polygon import Polygon
 from cartoframes import CartoDataFrame
 from cartoframes.auth import Credentials
 from cartoframes.data.clients.bigquery_client import BigQueryClient
+from cartoframes.data.observatory import Variable, Dataset
+from cartoframes.data.observatory.catalog.repository.entity_repo import EntityRepository
 from cartoframes.data.observatory.enrichment.enrichment_service import EnrichmentService, prepare_variables
-from cartoframes.data.observatory import Variable
+from cartoframes.exceptions import EnrichmentException
 from cartoframes.utils.geom_utils import to_geojson
 
 try:
@@ -151,8 +154,11 @@ class TestEnrichmentService(object):
 
         BigQueryClient._init_client = original
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Variable, 'get')
-    def test_prepare_variables(self, get_mock):
+    def test_prepare_variables(self, get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         variable_id = 'project.dataset.table.variable'
         variable = Variable({
             'id': variable_id,
@@ -182,3 +188,34 @@ class TestEnrichmentService(object):
             result = prepare_variables(case)
 
             assert result == [variable, variable]
+
+    @patch.object(EntityRepository, 'get_by_id')
+    @patch.object(Variable, 'get')
+    def test_prepare_variables_raises_if_not_available_in_bq(self, get_mock, entity_repo):
+        # mock dataset
+        entity_repo.return_value = Dataset({
+            'id': 'id',
+            'slug': 'slug',
+            'name': 'name',
+            'description': 'description',
+            'available_in': [],
+            'geography_id': 'geography'
+        })
+
+        variable = Variable({
+            'id': 'id',
+            'column_name': 'column',
+            'dataset_id': 'fake_name',
+            'slug': 'slug'
+        })
+
+        get_mock.return_value = variable
+
+        with pytest.raises(EnrichmentException) as e:
+            prepare_variables(variable)
+
+        error = """
+            The Dataset or the Geography of the Variable '{}' is not ready for Enrichment.
+            Please, contact us for more information.
+        """.format(variable.slug)
+        assert str(e.value) == error
diff --git a/tests/unit/data/observatory/enrichment/test_points_enrichment.py b/tests/unit/data/observatory/enrichment/test_points_enrichment.py
index 107a4fc8..d7704e5c 100644
--- a/tests/unit/data/observatory/enrichment/test_points_enrichment.py
+++ b/tests/unit/data/observatory/enrichment/test_points_enrichment.py
@@ -209,8 +209,11 @@ class TestPointsEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_points_with_filters(self, dataset_get_mock):
+    def test_enrichment_query_by_points_with_filters(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
diff --git a/tests/unit/data/observatory/enrichment/test_polygon_enrichment.py b/tests/unit/data/observatory/enrichment/test_polygon_enrichment.py
index c11b606c..bedee8d0 100644
--- a/tests/unit/data/observatory/enrichment/test_polygon_enrichment.py
+++ b/tests/unit/data/observatory/enrichment/test_polygon_enrichment.py
@@ -27,8 +27,11 @@ class TestPolygonEnrichment(object):
         self.credentials = None
         BigQueryClient._init_client = self.original_init_client
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_one_variable(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_one_variable(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -66,8 +69,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_two_variables(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_two_variables(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -113,8 +119,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_two_variables_different_tables(self, dataset_get_mock):
+    def test_enrichment_query_polygons_2_vars_different_tables(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -163,8 +172,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_two_variables_different_datasets(self, dataset_get_mock):
+    def test_enrichment_query_polygons_2_vars_different_datasets(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -214,8 +226,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_agg_empty_uses_variable_one(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_agg_empty_uses_variable_one(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -254,8 +269,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_agg_as_string(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_agg_as_string(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -294,8 +312,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_without_agg(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_without_agg(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -332,8 +353,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_agg_custom(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_agg_custom(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -386,8 +410,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_using_public_project(self, dataset_get_mock):
+    def test_enrichment_query_using_public_project(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'
@@ -423,8 +450,11 @@ class TestPolygonEnrichment(object):
 
         assert actual == expected
 
+    @patch('cartoframes.data.observatory.enrichment.enrichment_service._is_available_in_bq')
     @patch.object(Dataset, 'get')
-    def test_enrichment_query_by_polygons_with_filters(self, dataset_get_mock):
+    def test_enrichment_query_by_polygons_with_filters(self, dataset_get_mock, _is_available_in_bq_mock):
+        _is_available_in_bq_mock.return_value = True
+
         enrichment = Enrichment(credentials=self.credentials)
 
         temp_table_name = 'test_table'