fix quick test

PriorLabs · Jan 7, 2025 · 666c079 · 666c079
1 parent a8f1989
commit 666c079
Showing 1 changed file with 38 additions and 38 deletions.
diff --git a/tabpfn_client/tests/quick_test.py b/tabpfn_client/tests/quick_test.py
@@ -7,6 +7,7 @@
 """
 
 import logging
+from unittest.mock import patch
 
 from sklearn.datasets import load_breast_cancer, load_diabetes
 from sklearn.model_selection import train_test_split
@@ -18,41 +19,40 @@
 
 
 if __name__ == "__main__":
-    # set logging level to debug
-    # logging.basicConfig(level=logging.DEBUG)
-
-    use_server = True
-    # use_server = False
-
-    X, y = load_breast_cancer(return_X_y=True)
-    X_train, X_test, y_train, y_test = train_test_split(
-        X, y, test_size=0.33, random_state=42
-    )
-    y_train = y_train
-    y_test = y_test
-
-    tabpfn = TabPFNClassifier(n_estimators=3)
-    # print("checking estimator", check_estimator(tabpfn))
-    tabpfn.fit(X_train[:99], y_train[:99])
-    print("predicting")
-    print(tabpfn.predict(X_test))
-    print("predicting_proba")
-    print(tabpfn.predict_proba(X_test))
-
-    print(UserDataClient.get_data_summary())
-
-    X, y = load_diabetes(return_X_y=True)
-    X_train, X_test, y_train, y_test = train_test_split(
-        X, y, test_size=0.33, random_state=42
-    )
-
-    tabpfn = TabPFNRegressor(n_estimators=3)
-    # print("checking estimator", check_estimator(tabpfn))
-    tabpfn.fit(X_train[:99], y_train[:99])
-    print("predicting reg")
-    print(tabpfn.predict(X_test, output_type="mean"))
-
-    print(UserDataClient.get_data_summary())
-    # test predict_full
-    print("predicting ")
-    print(tabpfn.predict(X_test[:30], output_type="full", quantiles=[0.1, 0.5, 0.9]))
+    # Patch webbrowser.open to prevent browser login
+    with patch("webbrowser.open", return_value=False):
+        use_server = True
+        # use_server = False
+
+        X, y = load_breast_cancer(return_X_y=True)
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=0.33, random_state=42
+        )
+
+        tabpfn = TabPFNClassifier(n_estimators=3)
+        # print("checking estimator", check_estimator(tabpfn))
+        tabpfn.fit(X_train[:99], y_train[:99])
+        print("predicting")
+        print(tabpfn.predict(X_test))
+        print("predicting_proba")
+        print(tabpfn.predict_proba(X_test))
+
+        print(UserDataClient.get_data_summary())
+
+        X, y = load_diabetes(return_X_y=True)
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=0.33, random_state=42
+        )
+
+        tabpfn = TabPFNRegressor(n_estimators=3)
+        # print("checking estimator", check_estimator(tabpfn))
+        tabpfn.fit(X_train[:99], y_train[:99])
+        print("predicting reg")
+        print(tabpfn.predict(X_test, output_type="mean"))
+
+        print(UserDataClient.get_data_summary())
+        # test predict_full
+        print("predicting ")
+        print(
+            tabpfn.predict(X_test[:30], output_type="full", quantiles=[0.1, 0.5, 0.9])
+        )