Using print instead of log for displaying the device auth message (#170)

pykusto/_src/pyspark_client.py

@@ -6,7 +6,6 @@ import pandas as pd
 from azure.kusto.data import ClientRequestProperties, KustoClient
 
 from pykusto import PyKustoClient, NO_RETRIES, KustoResponse, KQL, RetryConfig
-from .logger import _logger
 
 
 class DataframeBasedKustoResponse(KustoResponse):
@@ -30,6 +29,7 @@ class PySparkKustoClient(PyKustoClient):
     """
     Handle to a Kusto cluster, to be used inside a PySpark notebook.
     """
+
     def __init__(self, cluster: str, linked_service: str = None, fetch_by_default: bool = True) -> None:
         """
         Create a new handle to a Kusto cluster. The value of "fetch_by_default" is used for current instance, and also passed on to database instances.
@@ -71,7 +71,7 @@ class PySparkKustoClient(PyKustoClient):
         assert self.__linked_service is None, "Device authentication can be used only when a linked_service was not provided to the client constructor"
         # noinspection PyProtectedMember
         device_auth = self.__spark_context._jvm.com.microsoft.kusto.spark.authentication.DeviceAuthentication(self.__cluster_name, "common")
-        _logger.info(device_auth.getDeviceCodeMessage())
+        print(device_auth.getDeviceCodeMessage())  # Logging is better than printing, but the PySpark notebook does not display logs by default
         self.option('accessToken', device_auth.acquireToken)
 
     # noinspection PyUnresolvedReferences,PyPackageRequirements

test/test_base.py

@@ -2,6 +2,8 @@ import json
 import logging
 import sys
 from concurrent.futures import Future
+from contextlib import contextmanager
+from io import StringIO
 from threading import Event
 from typing import Callable, Tuple, Any, List, Optional, Union, Type
 from unittest import TestCase
@@ -270,4 +272,15 @@ def nested_attribute_dict(attributes: str, value: Any) -> Any:
     return result
 
 
+@contextmanager
+def captured_stdout():
+    new_out = StringIO()
+    old_out = sys.stdout
+    try:
+        sys.stdout = new_out
+        yield sys.stdout
+    finally:
+        sys.stdout = old_out
+
+
 test_logger = logging.getLogger("pykusto_test")

test/test_pyspark_client.py

@@ -1,4 +1,3 @@
-import logging
 from unittest.mock import patch
 
 import pandas as pd
@@ -7,8 +6,7 @@ from pykusto import Query, PySparkKustoClient
 # noinspection PyProtectedMember
 from pykusto._src.expressions import _StringColumn, _NumberColumn
 # noinspection PyProtectedMember
-from pykusto._src.logger import _logger
+from test.test_base import TestBase, nested_attribute_dict, captured_stdout
-from test.test_base import TestBase, nested_attribute_dict
 
 
 class MockDataFrameReader:
@@ -112,11 +110,11 @@ class TestClient(TestBase):
         mock_spark_session = MockSparkSession(expected_df)
         mock_spark_context = MockSparkContext('MOCK_TOKEN')
 
-        with patch('pykusto._src.pyspark_client.PySparkKustoClient._PySparkKustoClient__get_spark_session_and_context', lambda s: (mock_spark_session, mock_spark_context)),\
+        with patch('pykusto._src.pyspark_client.PySparkKustoClient._PySparkKustoClient__get_spark_session_and_context', lambda s: (mock_spark_session, mock_spark_context)), \
-                self.assertLogs(_logger, logging.INFO) as cm:
+                captured_stdout() as out:
             client = PySparkKustoClient('https://help.kusto.windows.net/', fetch_by_default=False)
 
-        self.assertEqual(["INFO:pykusto:To sign in, use a lubricated goat to open the pod bay doors."], cm.output)
+        self.assertEqual("To sign in, use a lubricated goat to open the pod bay doors.", out.getvalue().strip())
         table = client['test_db']['mock_table']
         actual_df = Query(table).take(5).to_dataframe()
         self.assertTrue(expected_df.equals(actual_df))