From 7e357abcd93bd1852e4cac4e08c092fa76ea0bc6 Mon Sep 17 00:00:00 2001
From: Mauko Quiroga <mauko.quiroga@data.gouv.fr>
Date: Sat, 8 Feb 2020 18:28:52 +0100
Subject: [PATCH] Add test to get_memory_usage

---
 openfisca_core/data_storage.py                | 58 +++++++++++++------
 .../data_storage/test_in_memory_storage.py    | 13 +++++
 .../core/data_storage/test_on_disk_storage.py | 22 +++++++
 3 files changed, 76 insertions(+), 17 deletions(-)

diff --git a/openfisca_core/data_storage.py b/openfisca_core/data_storage.py
index 77e0cea6e5..8fcfcdc62a 100644
--- a/openfisca_core/data_storage.py
+++ b/openfisca_core/data_storage.py
@@ -25,25 +25,13 @@ def put(self, value: numpy.ndarray, period: periods.Period) -> None:
     def delete(self, period: Optional[periods.Period] = None) -> None:
         ...
 
+    @abc.abstractmethod
     def get_known_periods(self) -> List[periods.Period]:
         ...
 
-    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,
-            }
+    @abc.abstractmethod
+    def get_memory_usage(self) -> dict:
+        ...
 
 
 class InMemoryStorage(StorageLike):
@@ -97,6 +85,23 @@ def delete(self, period: Optional[periods.Period] = None) -> None:
     def get_known_periods(self) -> List[periods.Period]:
         return list(self._arrays.keys())
 
+    def get_memory_usage(self) -> dict:
+        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):
     """
@@ -175,9 +180,28 @@ def delete(self, period: Optional[periods.Period] = None) -> None:
     def get_known_periods(self) -> List[periods.Period]:
         return list(self._files.keys())
 
+    def get_memory_usage(self) -> dict:
+        if not self._files:
+            return {
+                "nb_files": 0,
+                "total_nb_bytes": 0,
+                "cell_size": numpy.nan,
+                }
+
+        nb_files = len(self._files)
+        file = next(iter(self._files.values()))
+        size = os.path.getsize(file)
+        array = self._decode_file(file)
+
+        return {
+            "nb_files": nb_files,
+            "total_nb_bytes": size * nb_files,
+            "cell_size": array.itemsize,
+            }
+
     def restore(self):
         self._files = files = {}
-        # Restore self._arrays from content of storage_dir.
+        # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):
             if not filename.endswith('.npy'):
                 continue
diff --git a/tests/core/data_storage/test_in_memory_storage.py b/tests/core/data_storage/test_in_memory_storage.py
index 239a9ebbae..3015a6b418 100644
--- a/tests/core/data_storage/test_in_memory_storage.py
+++ b/tests/core/data_storage/test_in_memory_storage.py
@@ -98,3 +98,16 @@ def test_get_known_periods(storage, value, period):
     result = storage.get_known_periods()
 
     assert result == [period]
+
+
+def test_get_memory_usage(storage, value, period):
+    storage = storage()
+    storage.put(value, period)
+
+    result = storage.get_memory_usage()
+
+    assert result == {
+        "nb_arrays": 1,
+        "total_nb_bytes": 8,
+        "cell_size": 8,
+        }
diff --git a/tests/core/data_storage/test_on_disk_storage.py b/tests/core/data_storage/test_on_disk_storage.py
index b852043c4e..125110f872 100644
--- a/tests/core/data_storage/test_on_disk_storage.py
+++ b/tests/core/data_storage/test_on_disk_storage.py
@@ -95,3 +95,25 @@ 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]
+
+
+def test_get_memory_usage(storage, value, period):
+    storage = storage()
+    storage.put(value, period)
+
+    result = storage.get_memory_usage()
+
+    assert result == {
+        "nb_files": 1,
+        "total_nb_bytes": 136,
+        "cell_size": 8,
+        }