diff --git a/chimedb/dataset/get.py b/chimedb/dataset/get.py index 42b8e03..9dada8d 100644 --- a/chimedb/dataset/get.py +++ b/chimedb/dataset/get.py @@ -105,14 +105,14 @@ def from_id(cls, state_id, load_data=True): try: _logger.debug(f"Loading state {state_id} from database...") if load_data: - _state_cache[state_id] = DatasetState.get( - DatasetState.id == state_id - ) + _state_cache[state_id] = DatasetState.get_by_id(state_id) else: - _state_cache[state_id] = DatasetState.select( - DatasetState.type, DatasetState.time - ).get() - except orm.DatasetState.DoesNotExist: + _state_cache[state_id] = ( + DatasetState.select(DatasetState.type, DatasetState.time) + .where(DatasetState.id == state_id) + .get() + ) + except DatasetState.DoesNotExist: _logger.warning(f"Could not find state {state_id}.") return None return _state_cache[state_id] @@ -191,8 +191,8 @@ def from_id(cls, ds_id): ) if ds_id not in _dataset_cache: try: - _dataset_cache[ds_id] = Dataset.get(Dataset.id == ds_id) - except orm.Dataset.DoesNotExist: + _dataset_cache[ds_id] = Dataset.get_by_id(ds_id) + except Dataset.DoesNotExist: _logger.warning(f"Could not find dataset {ds_id}.") return None return _dataset_cache[ds_id] diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 3b41a6a..b12fc1a 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -106,3 +106,17 @@ def tests(): # pre-fetch dget.index() tests() + + def test_get_missing(self): + # These don't exist + assert dget.Dataset.from_id("42") is None + assert dget.DatasetState.from_id("42") is None + + # But these do + ds = dget.Dataset.from_id("1337") + assert ds is not None + assert ds.id == "1337" + + state = dget.DatasetState.from_id("23") + assert state is not None + assert state.id == "23"