diff --git a/pykusto/_src/pyspark_client.py b/pykusto/_src/pyspark_client.py index 5377dff..e8c2f7b 100644 --- a/pykusto/_src/pyspark_client.py +++ b/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 diff --git a/test/test_base.py b/test/test_base.py index bb4feae..1a5bea0 100644 --- a/test/test_base.py +++ b/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") diff --git a/test/test_pyspark_client.py b/test/test_pyspark_client.py index 4c7bed2..88a2df0 100644 --- a/test/test_pyspark_client.py +++ b/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 +from test.test_base import TestBase, nested_attribute_dict, captured_stdout 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)),\ - self.assertLogs(_logger, logging.INFO) as cm: + with patch('pykusto._src.pyspark_client.PySparkKustoClient._PySparkKustoClient__get_spark_session_and_context', lambda s: (mock_spark_session, mock_spark_context)), \ + 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))