Add full warning assertions for Sample.from_frame id inference (#337)

Summary:
Added full warning assertions for `Sample.from_frame` id inference

Pull Request resolved: #337

Differential Revision: D93365012

Pulled By: talgalili

fbshipit-source-id: dd3327c85429a9882a89f0d580b15282db0f84aa

Author: neuralsorcerer

CHANGELOG.md

@@ -26,6 +26,11 @@
     to their first column before computation, matching validation behavior and
     returning scalar/Series outputs consistently.
 
+## Tests
+
+- **Expanded warning coverage for `Sample.from_frame()` ID inference**
+  - Added assertions that validate all three expected warnings are emitted when inferring an `id` column and default weights, including ID guessing, ID string casting, and automatic weight creation.
+
 # 0.16.0 (2026-02-09)
 
 ## New Features

tests/test_sample.py

@@ -174,17 +174,18 @@ def test_Sample_from_frame_id_column_detection(self) -> None:
         """
         # Test automatic id column detection
         df = pd.DataFrame({"id": (1, 2), "a": (1, 2)})
-        self.assertWarnsRegexp(
-            "Guessed id column name id for the data", Sample.from_frame, df
-        )
-        # TODO: add tests for the two other warnings:
-        # - self.assertWarnsRegexp("Casting id column to string", Sample.from_frame, df)
-        # - self.assertWarnsRegexp("No weights passed, setting all weights to 1", Sample.from_frame, df)
-        # Using the above would fail since the warnings are sent sequentially and using self.assertWarnsRegexp
-        # only catches the first warning.
-        self.assertEqual(
-            Sample.from_frame(df).id_column, pd.Series((1, 2), name="id").astype(str)
+        with self.assertLogs("balance", level="WARNING") as captured_logs:
+            sample = Sample.from_frame(df)
+
+        warning_messages = "\n".join(captured_logs.output)
+
+        self.assertIn("Guessed id column name id for the data", warning_messages)
+        self.assertIn("Casting id column to string", warning_messages)
+        self.assertIn(
+            "No weights passed. Adding a 'weight' column and setting all values to 1",
+            warning_messages,
         )
+        self.assertEqual(sample.id_column, pd.Series((1, 2), name="id").astype(str))
 
         # Test explicit id column specification
         df = pd.DataFrame({"b": (1, 2), "a": (1, 2)})