Avoid code repetition in test

Co-authored-by: Gonzalo <gonzaponte@gmail.com>
next-exp · Oct 28, 2021 · 30d55cd · 30d55cd
1 parent bef634c
commit 30d55cd
Showing 1 changed file with 19 additions and 36 deletions.
diff --git a/invisible_cities/database/load_db_test.py b/invisible_cities/database/load_db_test.py
@@ -193,23 +193,15 @@ def test_RadioactivityData_max_version():
 
     conn = sqlite3.connect(DB.get_db(db_file))
 
-    # test Activity table
-    query = '''SELECT * FROM Activity
-               WHERE G4Volume = '{0}' AND Isotope = '{1}'
-            '''
-    for (g4volume, isotope), req in activity.groupby(["G4Volume", "Isotope"]):
-        all_versions = pd.read_sql_query(query.format(g4volume, isotope), conn)
-        max_version = all_versions[all_versions.Version == all_versions.Version.max()]
-        np.testing.assert_allclose(req.TotalActivity, max_version.TotalActivity)
-
-    # test Efficiency table
-    query = '''SELECT * FROM Efficiency
-               WHERE G4Volume = '{0}' AND Isotope = '{1}'
-            '''
-    for (g4volume, isotope), req in efficiency.groupby(["G4Volume", "Isotope"]):
-        all_versions = pd.read_sql_query(query.format(g4volume, isotope), conn)
-        max_version = all_versions[all_versions.Version == all_versions.Version.max()]
-        np.testing.assert_allclose(req.MCEfficiency, max_version.MCEfficiency)
+    query = "SELECT * FROM {0} WHERE G4Volume = '{1}' AND Isotope = '{2}'"
+
+    for table, table_name, column in zip((      activity ,    efficiency ),
+                                         (     "Activity",   "Efficiency"),
+                                         ("TotalActivity", "MCEfficiency")):
+        for (g4volume, isotope), req in table.groupby(["G4Volume", "Isotope"]):
+            all_versions = pd.read_sql_query(query.format(table_name, g4volume, isotope), conn)
+            max_version  = all_versions[all_versions.Version == all_versions.Version.max()]
+            np.testing.assert_allclose(req[column], max_version[column])
 
 
 @mark.parametrize("version", (0, 1, 2))
@@ -221,22 +213,13 @@ def test_RadioactivityData_get_version(version):
 
     conn = sqlite3.connect(DB.get_db(db_file))
 
-    # test Activity table
-    query = '''SELECT * FROM Activity
-               WHERE G4Volume = '{0}' AND Isotope = '{1}'
-            '''
-    for (g4volume, isotope), req in activity.groupby(["G4Volume", "Isotope"]):
-        all_versions = pd.read_sql_query(query.format(g4volume, isotope), conn)
-        leq_version     = all_versions[all_versions.Version <= version]
-        closest_version = leq_version [ leq_version.Version == leq_version.Version.max()]
-        np.testing.assert_allclose(req.TotalActivity, closest_version.TotalActivity)
-
-    # test Efficiency table
-    query = '''SELECT * FROM Efficiency
-               WHERE G4Volume = '{0}' AND Isotope = '{1}'
-            '''
-    for (g4volume, isotope), req in efficiency.groupby(["G4Volume", "Isotope"]):
-        all_versions = pd.read_sql_query(query.format(g4volume, isotope), conn)
-        leq_version     = all_versions[all_versions.Version <= version]
-        closest_version = leq_version [ leq_version.Version == leq_version.Version.max()]
-        np.testing.assert_allclose(req.MCEfficiency, closest_version.MCEfficiency)
+    query = "SELECT * FROM {0} WHERE G4Volume = '{1}' AND Isotope = '{2}'"
+
+    for table, table_name, column in zip((      activity ,    efficiency ),
+                                         (     "Activity",   "Efficiency"),
+                                         ("TotalActivity", "MCEfficiency")):
+        for (g4volume, isotope), req in table.groupby(["G4Volume", "Isotope"]):
+            all_versions    = pd.read_sql_query(query.format(table_name, g4volume, isotope), conn)
+            leq_version     = all_versions[all_versions.Version <= version]
+            closest_version = leq_version [ leq_version.Version == leq_version.Version.max()]
+            np.testing.assert_allclose(req[column], closest_version[column])