test: replace fixtures with dynamic mock data

Author: mplatzer

tests/end_to_end/fixtures/baseball/players-holdout.parquet

@@ -12,43 +12,55 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import uuid
 from pathlib import Path
 
 import pandas as pd
+import numpy as np
 
-from mostlyai.qa.report_from_statistics import report_from_statistics
-from mostlyai.qa.report import report
+from mostlyai import qa
 
-baseball_path = Path(os.path.realpath(__file__)).parent / "fixtures" / "baseball"
-census_path = Path(os.path.realpath(__file__)).parent / "fixtures" / "census"
+
+def mock_data(n):
+    df = pd.DataFrame(
+        {
+            "int": pd.Series(np.random.choice(list(range(10)) + [np.nan], n), dtype="Int64[pyarrow]"),
+            "float": pd.Series(
+                np.random.choice(list(np.random.uniform(size=10)) + [np.nan], n), dtype="float64[pyarrow]"
+            ),
+            "cat": pd.Series(np.random.choice(["f", "m", np.nan], n), dtype="string[pyarrow]"),
+            "bool": pd.Series(np.random.choice([True, False, np.nan], n), dtype="bool[pyarrow]"),
+            "date": pd.Series(np.random.choice([np.datetime64("today", "D"), np.nan], n), dtype="datetime64[ns]"),
+            "text": pd.Series([str(uuid.uuid4())[:4] for _ in range(n)], dtype="object"),
+        }
+    )
+    return df
 
 
 def test_report_flat(tmp_path):
     statistics_path = tmp_path / "statistics"
-    syn_tgt_data = pd.read_parquet(census_path / "census-synthetic.parquet")
-    trn_tgt_data = pd.read_parquet(census_path / "census-training.parquet")
-    hol_tgt_data = pd.read_parquet(census_path / "census-holdout.parquet")
-    report_path, metrics = report(
+    syn_tgt_data = mock_data(220)
+    trn_tgt_data = mock_data(180)
+    hol_tgt_data = mock_data(140)
+    report_path, metrics = qa.report(
         syn_tgt_data=syn_tgt_data,
         trn_tgt_data=trn_tgt_data,
         hol_tgt_data=hol_tgt_data,
         statistics_path=statistics_path,
-        max_sample_size_accuracy=2000,
-        max_sample_size_embeddings=200,
+        max_sample_size_accuracy=120,
+        max_sample_size_embeddings=80,
     )
 
     assert report_path.exists()
 
     accuracy = metrics.accuracy
-    assert 0.5 <= accuracy.overall <= 1.0
-    assert 0.5 <= accuracy.univariate <= 1.0
-    assert 0.5 <= accuracy.bivariate <= 1.0
+    assert 0.3 <= accuracy.overall <= 1.0
+    assert 0.3 <= accuracy.univariate <= 1.0
+    assert 0.3 <= accuracy.bivariate <= 1.0
     assert accuracy.coherence is None
-    assert 0.8 <= accuracy.overall_max <= 1.0
-    assert 0.8 <= accuracy.univariate_max <= 1.0
-    assert 0.8 <= accuracy.bivariate_max <= 1.0
+    assert 0.3 <= accuracy.overall_max <= 1.0
+    assert 0.3 <= accuracy.univariate_max <= 1.0
+    assert 0.3 <= accuracy.bivariate_max <= 1.0
 
     similarity = metrics.similarity
     assert 0.8 <= similarity.cosine_similarity_training_synthetic <= 1.0
@@ -63,11 +75,11 @@ def test_report_flat(tmp_path):
     assert 0 <= distances.dcr_holdout <= 1.0
     assert 0 <= distances.dcr_share <= 1.0
 
-    report_path = report_from_statistics(
+    report_path = qa.report_from_statistics(
         syn_tgt_data=syn_tgt_data,
         statistics_path=statistics_path,
-        max_sample_size_accuracy=2000,
-        max_sample_size_embeddings=200,
+        max_sample_size_accuracy=110,
+        max_sample_size_embeddings=70,
     )
 
     assert report_path.exists()
@@ -77,44 +89,51 @@ def test_report_sequential(tmp_path):
     statistics_path = tmp_path / "statistics"
     report_path = Path(tmp_path / "my-report.html")
 
-    trn_tgt_data = pd.read_parquet(baseball_path / "seasons-training.parquet")
-    hol_tgt_data = pd.read_parquet(baseball_path / "seasons-holdout.parquet")
-    syn_tgt_data = pd.read_parquet(baseball_path / "seasons-synthetic.parquet")
-    trn_ctx_data = pd.read_parquet(baseball_path / "players-training.parquet")
-    hol_ctx_data = pd.read_parquet(baseball_path / "players-holdout.parquet")
-    syn_ctx_data = pd.concat([trn_ctx_data, hol_ctx_data])
-
-    report_path, metrics = report(
+    # generate mock context data
+    syn_ctx_data = mock_data(220).reset_index(names="id")
+    trn_ctx_data = mock_data(180).reset_index(names="id")
+    hol_ctx_data = mock_data(140).reset_index(names="id")
+
+    # generate mock sequential target data
+    syn_tgt_data = mock_data(220 * 3)
+    syn_tgt_data["ctx_id"] = np.random.choice(syn_ctx_data["id"], 220 * 3)
+    trn_tgt_data = mock_data(180 * 4)
+    trn_tgt_data["ctx_id"] = np.random.choice(trn_ctx_data["id"], 180 * 4)
+    hol_tgt_data = mock_data(140 * 4)
+    hol_tgt_data["ctx_id"] = np.random.choice(hol_ctx_data["id"], 140 * 4)
+
+    # generate report
+    report_path, metrics = qa.report(
         syn_tgt_data=syn_tgt_data,
         trn_tgt_data=trn_tgt_data,
         hol_tgt_data=hol_tgt_data,
         syn_ctx_data=syn_ctx_data,
         trn_ctx_data=trn_ctx_data,
         hol_ctx_data=hol_ctx_data,
         ctx_primary_key="id",
-        tgt_context_key="players_id",
+        tgt_context_key="ctx_id",
         report_path=report_path,
         statistics_path=statistics_path,
-        max_sample_size_accuracy=2000,
-        max_sample_size_embeddings=200,
+        max_sample_size_accuracy=120,
+        max_sample_size_embeddings=80,
     )
 
     assert report_path.exists()
 
     accuracy = metrics.accuracy
-    assert 0.8 <= accuracy.overall <= 1.0
-    assert 0.8 <= accuracy.univariate <= 1.0
-    assert 0.8 <= accuracy.bivariate <= 1.0
-    assert 0.8 <= accuracy.coherence <= 1.0
-    assert 0.8 <= accuracy.overall_max <= 1.0
-    assert 0.8 <= accuracy.univariate_max <= 1.0
-    assert 0.8 <= accuracy.bivariate_max <= 1.0
-    assert 0.8 <= accuracy.coherence_max <= 1.0
+    assert 0.3 <= accuracy.overall <= 1.0
+    assert 0.3 <= accuracy.univariate <= 1.0
+    assert 0.3 <= accuracy.bivariate <= 1.0
+    assert 0.3 <= accuracy.coherence <= 1.0
+    assert 0.3 <= accuracy.overall_max <= 1.0
+    assert 0.3 <= accuracy.univariate_max <= 1.0
+    assert 0.3 <= accuracy.bivariate_max <= 1.0
+    assert 0.3 <= accuracy.coherence_max <= 1.0
 
     similarity = metrics.similarity
-    assert 0.8 <= similarity.cosine_similarity_training_synthetic <= 1.0
+    assert 0.3 <= similarity.cosine_similarity_training_synthetic <= 1.0
     assert 0.0 <= similarity.discriminator_auc_training_synthetic <= 1.0
-    assert 0.8 <= similarity.cosine_similarity_training_holdout <= 1.0
+    assert 0.3 <= similarity.cosine_similarity_training_holdout <= 1.0
     assert 0.0 <= similarity.discriminator_auc_training_holdout <= 1.0
 
     distances = metrics.distances
@@ -124,13 +143,13 @@ def test_report_sequential(tmp_path):
     assert 0 <= distances.dcr_holdout <= 1.0
     assert 0 <= distances.dcr_share <= 1.0
 
-    report_path = report_from_statistics(
+    report_path = qa.report_from_statistics(
         syn_tgt_data=syn_tgt_data,
         syn_ctx_data=syn_ctx_data,
         ctx_primary_key="id",
-        tgt_context_key="players_id",
-        max_sample_size_accuracy=3000,
-        max_sample_size_embeddings=300,
+        tgt_context_key="ctx_id",
+        max_sample_size_accuracy=130,
+        max_sample_size_embeddings=90,
         statistics_path=statistics_path,
     )
 
@@ -144,7 +163,7 @@ def test_report_flat_rare(tmp_path):
     syn_tgt_data = pd.DataFrame({"x": ["_RARE_" for _ in range(100)]})
     trn_tgt_data = pd.DataFrame({"x": [str(uuid.uuid4()) for _ in range(100)]})
     hol_tgt_data = pd.DataFrame({"x": [str(uuid.uuid4()) for _ in range(100)]})
-    _, metrics = report(
+    _, metrics = qa.report(
         syn_tgt_data=syn_tgt_data,
         trn_tgt_data=trn_tgt_data,
         hol_tgt_data=hol_tgt_data,
@@ -155,7 +174,7 @@ def test_report_flat_rare(tmp_path):
 
     # test case where rare values are not protected, and we leak trn into synthetic
     syn_tgt_data = pd.DataFrame({"x": trn_tgt_data["x"].sample(100, replace=True)})
-    _, metrics = report(
+    _, metrics = qa.report(
         syn_tgt_data=syn_tgt_data,
         trn_tgt_data=trn_tgt_data,
         hol_tgt_data=hol_tgt_data,
@@ -168,7 +187,7 @@ def test_report_flat_rare(tmp_path):
 def test_report_flat_early_exit(tmp_path):
     # test early exit for dfs with <100 rows
     df = pd.DataFrame({"col": list(range(99))})
-    _, metrics = report(syn_tgt_data=df, trn_tgt_data=df, hol_tgt_data=df)
+    _, metrics = qa.report(syn_tgt_data=df, trn_tgt_data=df, hol_tgt_data=df)
     assert metrics is None
 
 
@@ -196,7 +215,7 @@ def make_dfs(
         syn_ctx_data = trn_ctx_data = hol_ctx_data = ctx_df
         syn_tgt_data = trn_tgt_data = hol_tgt_data = tgt_df
         early_term = df_dict.pop("early_term")
-        _, metrics = report(
+        _, metrics = qa.report(
             syn_tgt_data=syn_tgt_data,
             trn_tgt_data=trn_tgt_data,
             hol_tgt_data=hol_tgt_data,
@@ -211,7 +230,7 @@ def make_dfs(
 
 def test_report_few_holdout_records(tmp_path):
     tgt = pd.DataFrame({"id": list(range(100)), "col": ["a"] * 100})
-    _, metrics = report(
+    _, metrics = qa.report(
         syn_tgt_data=tgt,
         trn_tgt_data=tgt,
         hol_tgt_data=tgt[:10],
@@ -223,7 +242,7 @@ def test_report_sequential_few_records(tmp_path):
     # ensure that we don't crash in case of dominant zero-seq-length
     ctx = pd.DataFrame({"id": list(range(1000))})
     tgt = pd.DataFrame({"id": [1, 2, 3, 4, 5] * 100, "col": ["a"] * 500})
-    _, metrics = report(
+    _, metrics = qa.report(
         syn_tgt_data=tgt,
         trn_tgt_data=tgt,
         hol_tgt_data=tgt,
@@ -245,14 +264,30 @@ def test_odd_column_names(tmp_path):
             "3": values,
         }
     )
-    path, metrics = report(
+    path, metrics = qa.report(
         syn_tgt_data=df,
         trn_tgt_data=df,
         statistics_path=tmp_path / "stats",
     )
     assert metrics is not None
-    path = report_from_statistics(
+    path = qa.report_from_statistics(
         syn_tgt_data=df,
         statistics_path=tmp_path / "stats",
     )
     assert path is not None
+
+
+def test_missing(tmp_path):
+    df1 = mock_data(100)
+    df2 = df1.copy()
+    df2.loc[:, :] = np.nan
+    _, metrics = qa.report(
+        syn_tgt_data=df1,
+        trn_tgt_data=df2,
+    )
+    assert metrics is not None
+    _, metrics = qa.report(
+        syn_tgt_data=df2,
+        trn_tgt_data=df1,
+    )
+    assert metrics is not None