Add test case for generated experiments

examples/tracking_merlion.py

@@ -28,8 +28,8 @@
 - https://github.com/numenta/NAB
 """
 
+import datetime as dt
 import os
-from datetime import timedelta
 
 import mlflow
 import numpy as np
@@ -155,7 +155,7 @@ def save_anomalies(table_name: str, anomalies: pd.DataFrame, mttd: int):
             rows.append(
                 (
                     index,
-                    index + timedelta(seconds=mttd),
+                    index + dt.timedelta(seconds=mttd),
                     row["anom_score"],
                     row["experiment_id"],
                     row["run_id"],

tests/test_examples.py

@@ -100,6 +100,23 @@ def test_tracking_merlion(reset_database, engine: sa.Engine, tracking_store: Sql
 
     # Verify database content.
     with tracking_store.ManagedSessionMaker() as session:
-        assert session.query(SqlExperiment).count() == 2
         assert session.query(SqlMetric).count() == 4
         assert session.query(SqlParam).count() == 5
+        experiments = session.query(SqlExperiment)
+        assert experiments.count() == 2  # default experiment + the current run = 2
+
+        experiment_id = experiments.first().experiment_id
+
+    with engine.begin() as conn:
+        conn.execute(sa.sql.text("REFRESH TABLE doc.anomalies"))
+
+        # Anomalies that have been manually annotated
+        anomalies_manual = conn.execute(sa.sql.text("SELECT * FROM doc.anomalies WHERE experiment_id IS NULL"))
+        assert anomalies_manual.rowcount == 3
+
+        # Anomalies that have been generated by the model
+        anomalies_generated = conn.execute(
+            sa.sql.text("SELECT * FROM doc.anomalies WHERE experiment_id = :experiment_id"),
+            parameters={"experiment_id": experiment_id},
+        )
+        assert anomalies_generated.rowcount == 3