From 1c663bc9522f1ef5ae8e571996264af616f82b0a Mon Sep 17 00:00:00 2001
From: judynah <mrs.nataliawalczak@gmail.com>
Date: Mon, 25 Nov 2024 12:14:59 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20Update=20docs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/viadot/sources/tm1.py | 17 ++++++++---------
 tests/unit/test_tm1.py    |  2 +-
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/viadot/sources/tm1.py b/src/viadot/sources/tm1.py
index 153b6e79d..72639822a 100644
--- a/src/viadot/sources/tm1.py
+++ b/src/viadot/sources/tm1.py
@@ -9,6 +9,7 @@
 from viadot.config import get_source_credentials
 from viadot.exceptions import ValidationError
 from viadot.sources.base import Source
+from viadot.utils import add_viadot_metadata_columns
 
 
 class TM1Credentials(BaseModel):
@@ -28,7 +29,7 @@ class TM1Credentials(BaseModel):
 
 
 class TM1(Source):
-    """Class for downloading data from TM1 Software using TM1py library."""
+    """TM1 connector."""
 
     def __init__(
         self,
@@ -45,17 +46,14 @@ def __init__(
         *args,
         **kwargs,
     ):
-        """Creating an instance of TM1 source class.
-
-            To download the data to the dataframe user needs to specify MDX query or
-            combination of cube and view.
+        """Create a TM1 connector instance.
 
         Args:
             credentials (dict[str, Any], optional): Credentials stored in a dictionary.
                 Required credentials: username, password, address, port.
                 Defaults to None.
-            config_key (str, optional): Credential key to dictionary where credentials
-                are stored. Defaults to "TM1".
+            config_key (str, optional): The key in the viadot config holding relevant
+                credentials. Defaults to "TM1".
             mdx_query (str, optional): MDX select query needed to download the data.
                 Defaults to None.
             cube (str, optional): Cube name from which data will be downloaded.
@@ -169,8 +167,9 @@ def get_available_elements(self) -> list:
             dimension_name=self.dimension, hierarchy_name=self.hierarchy
         )
 
+    @add_viadot_metadata_columns
     def to_df(self, if_empty: Literal["warn", "fail", "skip"] = "skip") -> pd.DataFrame:
-        """Function for downloading data from TM1 to pd.DataFrame.
+        """Download data into a pandas DataFrame.
 
             To download the data to the dataframe user needs to specify MDX query
             or combination of cube and view.
@@ -180,7 +179,7 @@ def to_df(self, if_empty: Literal["warn", "fail", "skip"] = "skip") -> pd.DataFr
             DataFrame is empty. Defaults to "skip".
 
         Returns:
-            pd.DataFrame: DataFrame with data downloaded from TM1 view.
+            pd.DataFrame: DataFrame with the data.
 
         Raises:
             ValidationError: When mdx and cube + view are not specified
diff --git a/tests/unit/test_tm1.py b/tests/unit/test_tm1.py
index 0569cc5e2..3e4906316 100644
--- a/tests/unit/test_tm1.py
+++ b/tests/unit/test_tm1.py
@@ -98,7 +98,7 @@ def test_to_df(tm1):
         result = tm1.to_df()
         assert isinstance(result, pd.DataFrame)
         assert not result.empty
-        assert len(result.columns) == 2
+        assert len(result.columns) == 4
 
 def test_to_df_fail(tm1_mock):
     with pytest.raises(ValidationError) as excinfo: