Add test to get_known_periods

openfisca_core/data_storage.py

@@ -2,7 +2,7 @@
 import collections
 import os
 import shutil
-from typing import Optional
+from typing import List, Optional
 
 import numpy
 
@@ -14,6 +14,8 @@
 
 class StorageLike(Protocol):
 
+    _arrays: dict
+
     @abc.abstractmethod
     def get(self, period: periods.Period) -> Optional[numpy.ndarray]:
         ...
@@ -29,6 +31,26 @@ def delete(
             ) -> Optional[numpy.ndarray]:
         ...
 
+    def get_known_periods(self) -> List[periods.Period]:
+        return list(self._arrays.keys())
+
+    def get_memory_usage(self):
+        if not self._arrays:
+            return {
+                "nb_arrays": 0,
+                "total_nb_bytes": 0,
+                "cell_size": numpy.nan,
+                }
+
+        nb_arrays = len(self._arrays)
+        array = next(iter(self._arrays.values()))
+
+        return {
+            "nb_arrays": nb_arrays,
+            "total_nb_bytes": array.nbytes * nb_arrays,
+            "cell_size": array.itemsize,
+            }
+
 
 class InMemoryStorage(StorageLike):
     """
@@ -77,33 +99,13 @@ def delete(
 
         return self._arrays.pop(period)
 
-    def get_known_periods(self):
-        return self._arrays.keys()
-
-    def get_memory_usage(self):
-        if not self._arrays:
-            return {
-                "nb_arrays": 0,
-                "total_nb_bytes": 0,
-                "cell_size": numpy.nan,
-                }
-
-        nb_arrays = len(self._arrays)
-        array = next(iter(self._arrays.values()))
-
-        return {
-            "nb_arrays": nb_arrays,
-            "total_nb_bytes": array.nbytes * nb_arrays,
-            "cell_size": array.itemsize,
-            }
-
 
 class OnDiskStorage(StorageLike):
     """
     Low-level class responsible for storing and retrieving calculated vectors on disk.
     """
 
-    _files: dict
+    _arrays: dict
     _enums: dict
     is_eternal: bool
     preserve_storage_dir: bool
@@ -115,7 +117,7 @@ def __init__(
             is_eternal: bool = False,
             preserve_storage_dir: bool = False,
             ) -> None:
-        self._files = {}
+        self._arrays = {}
         self._enums = {}
         self.is_eternal = is_eternal
         self.preserve_storage_dir = preserve_storage_dir
@@ -133,7 +135,7 @@ def get(self, period: periods.Period) -> Optional[numpy.ndarray]:
             period = periods.period(periods.ETERNITY)
 
         try:
-            values = self._files[period]
+            values = self._arrays[period]
             return self._decode_file(values)
 
         except KeyError:
@@ -153,29 +155,26 @@ def put(self, value: numpy.ndarray, period: periods.Period) -> None:
             value = value.view(numpy.ndarray)
 
         numpy.save(path, value)
-        self._files[period] = path
+        self._arrays[period] = path
 
     def delete(
             self,
             period: Optional[periods.Period] = None,
             ) -> Optional[numpy.ndarray]:
         if period is None:
-            self._files = {}
+            self._arrays = {}
             return None
 
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
 
         period = periods.period(period)
 
-        return self._files.pop(period)
-
-    def get_known_periods(self):
-        return self._files.keys()
+        return self._arrays.pop(period)
 
     def restore(self):
-        self._files = files = {}
-        # Restore self._files from content of storage_dir.
+        self._arrays = files = {}
+        # Restore self._arrays from content of storage_dir.
         for filename in os.listdir(self.storage_dir):
             if not filename.endswith('.npy'):
                 continue

tests/core/data_storage/test_in_memory_storage.py

@@ -89,3 +89,12 @@ def test_delete_when_is_eternal(eternal_storage, value):
     result = storage.get("qwerty"), storage.get("azerty")
 
     assert result == (None, None)
+
+
+def test_get_known_periods(storage, value, period):
+    storage = storage()
+    storage.put(value, period)
+
+    result = storage.get_known_periods()
+
+    assert result == [period]