diff --git a/django_project/spw/generator/main.py b/django_project/spw/generator/main.py index fff1b23..13414a7 100644 --- a/django_project/spw/generator/main.py +++ b/django_project/spw/generator/main.py @@ -15,7 +15,7 @@ from django.contrib.gis.geos import Point, Polygon from django.utils import timezone -from gap.models import Dataset, DatasetAttribute, CastType +from gap.models import Dataset, DatasetAttribute, CastType, DatasetStore from gap.providers import TomorrowIODatasetReader, TIO_PROVIDER from gap.utils.reader import DatasetReaderInput from spw.models import RModel, RModelExecutionLog, RModelExecutionStatus @@ -77,7 +77,8 @@ def calculate_from_point_attrs(): """Return attributes that are being used in calculate from point.""" return DatasetAttribute.objects.filter( attribute__variable_name__in=ATTRIBUTES, - dataset__provider__name=TIO_PROVIDER + dataset__provider__name=TIO_PROVIDER, + dataset__store_type=DatasetStore.EXT_API ) diff --git a/django_project/spw/tests/test_generator.py b/django_project/spw/tests/test_generator.py index c791247..18b3e36 100644 --- a/django_project/spw/tests/test_generator.py +++ b/django_project/spw/tests/test_generator.py @@ -29,7 +29,9 @@ SPWOutput, _fetch_timelines_data, _fetch_ltn_data, - calculate_from_point + calculate_from_point, + calculate_from_point_attrs, + VAR_MAPPING_REVERSE ) from spw.models import RModelExecutionLog, RModelExecutionStatus @@ -253,3 +255,23 @@ def test_calculate_from_point( self.assertTrue(log.input_file) self.assertTrue(log.output) self.assertEqual(log.status, RModelExecutionStatus.SUCCESS) + + +class TestSPWAttrs(TestCase): + """Test fetch attributes for SPW Tio.""" + + fixtures = [ + '2.provider.json', + '3.station_type.json', + '4.dataset_type.json', + '5.dataset.json', + '6.unit.json', + '7.attribute.json', + '8.dataset_attribute.json' + ] + + def test_attrs(self): + """Test to ensure attrs is correct for Tio API Call.""" + attrs = calculate_from_point_attrs() + for k, v in VAR_MAPPING_REVERSE.items(): + self.assertTrue(attrs.filter(source=k).exists())