Add tests for the multi-table synthesizer code (#69)

    Adding multi-table synthesizer tests
    Better way of checking if a test is being run in github actions by checking an environment variable that only exists there
    A few fixes to the current synthesizer code
    Fixing the test data to include a missing table relation that is important for the synthesizer, and re-generating and uploading the assertion data from github because it affects the results of existing tests

Author: lotif

.github/workflows/integration_tests.yml

@@ -58,6 +58,10 @@ jobs:
       - name: Install the project
         run: uv sync --all-extras --dev
 
+      - name: Is running on CI environment (GitHub Actions)?
+        run: |
+          python -c "import os; print('Result: ', os.getenv('GITHUB_ACTIONS', 'Not set'))"
+
       - name: Install dependencies and check code
         run: |
-          uv run pytest -m "integration_test"
+          uv run pytest -m "integration_test" --log-cli-level=WARNING

.github/workflows/unit_tests.yml

@@ -60,4 +60,4 @@ jobs:
 
       - name: Install dependencies and check code
         run: |
-          uv run pytest -m "not integration_test"
+          uv run pytest -m "not integration_test" --log-cli-level=WARNING

src/midst_toolkit/models/clavaddpm/synthesizer.py

@@ -786,7 +786,10 @@ def clava_synthesizing(  # noqa: PLR0915, PLR0912
                 "df": child_final_df,
                 "keys": child_primary_keys_arr.flatten().tolist(),
             }
-        with open(os.path.join(save_dir, "before_matching/synthetic_tables.pkl"), "wb") as file:
+
+        before_matching_dir = save_dir / "before_matching"
+        before_matching_dir.mkdir(parents=True, exist_ok=True)
+        with open(before_matching_dir / "synthetic_tables.pkl", "wb") as file:
             pickle.dump(synthetic_tables, file)
 
     synthesizing_end_time = time.time()
@@ -800,12 +803,8 @@ def clava_synthesizing(  # noqa: PLR0915, PLR0912
 
     cleaned_tables: dict[str, pd.DataFrame] = {}
     for table_key, table_val in final_tables.items():
-        if "account_id" in tables[table_key]["original_cols"]:
-            cols = tables[table_key]["original_cols"]
-            cols.remove("account_id")
-        else:
-            cols = tables[table_key]["original_cols"]
-        cleaned_tables[table_key] = pd.DataFrame(table_val[cols])
+        column_names = [column_name for column_name in tables[table_key]["original_cols"] if "_id" not in column_name]
+        cleaned_tables[table_key] = pd.DataFrame(table_val[column_names])
 
     for cleaned_key, cleaned_val in cleaned_tables.items():
         table_dir = os.path.join(

tests/integration/assets/multi_table/assertion_data/cleaned_tables.pkl

@@ -0,0 +1,51 @@
+{
+	"X_gen": [
+		[
+			-89.64320001648818,
+			2.329031172032132,
+			-122.97271923749832,
+            552.3706152861826,
+            353.47951217405426,
+            -63.164915493559306,
+            -42.27259013378604,
+            244.21392290993887
+        ],
+		[-0.4694302555020733,
+            15.336690906361277,
+            -48.59970780139716,
+            -358.65097509895173,
+            411.39200743280094,
+            415.9651477725036,
+            -12.980662539762594,
+            -370.11192775534397
+        ],
+		[5.009930498133295,
+            -220.79264470424582,
+            -4.129379545636459,
+            -188.011555249935,
+            218.10979023918082,
+            221.16927688555808,
+            49.89701474616661,
+            -194.37953943919408
+        ],
+		[-1.3109146973467711,
+            -73.2679936874503,
+            -9.218660554989645,
+            -389.99286808084486,
+            -490.3925197112697,
+            -423.00630661809424,
+            432.9884292987812,
+            -397.6777786014056
+        ],
+		[8.342572127948289,
+            9.36842404400312,
+            -72.28739585181947,
+            -489.4411862829012,
+            563.4325829362252,
+            568.3398615720979,
+            -16.894123940346486,
+            -504.9528775096839
+        ]
+	],
+	"y_gen": [1, 0, 1, 0, 0]
+}

tests/integration/assets/multi_table/assertion_data/conditional_samples.pt

@@ -1,5 +1,6 @@
 {
     "relation_order": [
+        [null, "account"],
         ["account", "trans"]
     ],
     "tables": {