test: remove boilerplate code

ing-bank · Nov 24, 2021 · 871d2c6 · 871d2c6
1 parent 8efc072
commit 871d2c6
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/tests/popmon/base/test_module.py b/tests/popmon/base/test_module.py
@@ -5,23 +5,22 @@
 
 def test_popmon_module():
     class Scaler(Module):
+        _input_keys = ("input_key",)
+        _output_keys = ("output_key",)
+
         def __init__(self, input_key, output_key, mean, std):
             super().__init__()
             self.input_key = input_key
             self.output_key = output_key
             self.mean = mean
             self.std = std
 
-        def transform(self, datastore):
-            input_array = self.get_datastore_object(
-                datastore, self.input_key, dtype=np.ndarray
-            )
+        def transform(self, input_array: np.ndarray):
             res = input_array - np.mean(input_array)
             res = res / np.std(res)
             res = res * self.std
             res = res + self.mean
-            datastore[self.output_key] = res
-            return datastore
+            return res
 
     test_module = Scaler(input_key="x", output_key="scaled_x", mean=2.0, std=0.3)