refactor(toggl): extract Toggl DataFrame init

Author: thekaveman

compiler_admin/services/toggl.py

@@ -49,6 +49,28 @@ def _get_last_name(email: str):
     return last_name
 
 
+def _prepare_input(source_path: str | TextIO, column_renames: dict = {}) -> pd.DataFrame:
+    """Parse and prepare CSV data from `source_path` into an initial `pandas.DataFrame`."""
+    df = files.read_csv(source_path, usecols=TOGGL_COLUMNS, parse_dates=["Start date"], cache_dates=True)
+
+    df["Start time"] = df["Start time"].apply(_str_timedelta)
+    df["Duration"] = df["Duration"].apply(_str_timedelta)
+
+    # assign First and Last name
+    df["First name"] = df["Email"].apply(_get_first_name)
+    df["Last name"] = df["Email"].apply(_get_last_name)
+
+    # calculate hours as a decimal from duration timedelta
+    df["Hours"] = (df["Duration"].dt.total_seconds() / 3600).round(2)
+
+    df.sort_values(["Start date", "Start time", "Email"], inplace=True)
+
+    if column_renames:
+        df.rename(columns=column_renames, inplace=True)
+
+    return df
+
+
 def _str_timedelta(td: str):
     """Convert a string formatted duration (e.g. 01:30) to a timedelta."""
     return pd.to_timedelta(pd.to_datetime(td, format="%H:%M:%S").strftime("%H:%M:%S"))
@@ -78,14 +100,9 @@ def convert_to_harvest(
     if client_name is None:
         client_name = os.environ.get("HARVEST_CLIENT_NAME")
 
-    # read CSV file, parsing dates and times
-    source = files.read_csv(source_path, usecols=TOGGL_COLUMNS, parse_dates=["Start date"], cache_dates=True)
-    source["Start time"] = source["Start time"].apply(_str_timedelta)
-    source["Duration"] = source["Duration"].apply(_str_timedelta)
-    source.sort_values(["Start date", "Start time", "Email"], inplace=True)
-
-    # rename columns that can be imported as-is
-    source.rename(columns={"Project": "Project", "Description": "Notes", "Start date": "Date"}, inplace=True)
+    source = _prepare_input(
+        source_path=source_path, column_renames={"Project": "Project", "Description": "Notes", "Start date": "Date"}
+    )
 
     # update static calculated columns
     source["Client"] = client_name
@@ -95,13 +112,6 @@ def convert_to_harvest(
     project_info = files.JsonFileCache("TOGGL_PROJECT_INFO")
     source["Project"] = source["Project"].apply(lambda x: project_info.get(key=x, default=x))
 
-    # assign First and Last name
-    source["First name"] = source["Email"].apply(_get_first_name)
-    source["Last name"] = source["Email"].apply(_get_last_name)
-
-    # calculate hours as a decimal from duration timedelta
-    source["Hours"] = (source["Duration"].dt.total_seconds() / 3600).round(2)
-
     files.write_csv(output_path, source, columns=output_cols)
 
 

tests/services/test_toggl.py

@@ -10,14 +10,15 @@
 import compiler_admin.services.toggl
 from compiler_admin.services.toggl import (
     __name__ as MODULE,
-    files,
-    TOGGL_COLUMNS,
-    HARVEST_COLUMNS,
     _get_first_name,
     _get_last_name,
+    _prepare_input,
     _str_timedelta,
     convert_to_harvest,
     download_time_entries,
+    TOGGL_COLUMNS,
+    HARVEST_COLUMNS,
+    files,
 )
 
 
@@ -137,10 +138,9 @@ def test_str_timedelta():
     assert result.total_seconds() == (1 * 60 * 60) + (30 * 60) + 15
 
 
-def test_convert_to_harvest_mocked(toggl_file, spy_files, mock_google_user_info):
-    mock_google_user_info.return_value = {}
-
-    convert_to_harvest(toggl_file, client_name=None)
+@pytest.mark.usefixtures("mock_google_user_info")
+def test_prepare_input(toggl_file, spy_files):
+    df = _prepare_input(toggl_file)
 
     spy_files.read_csv.assert_called_once()
     call_args = spy_files.read_csv.call_args
@@ -149,6 +149,31 @@ def test_convert_to_harvest_mocked(toggl_file, spy_files, mock_google_user_info)
     assert call_args.kwargs["parse_dates"] == ["Start date"]
     assert call_args.kwargs["cache_dates"] is True
 
+    df_cols = df.columns.to_list()
+    assert set(df_cols) <= set(TOGGL_COLUMNS) or set(TOGGL_COLUMNS) <= set(df_cols)
+
+    assert "First name" in df_cols
+    assert "Last name" in df_cols
+    assert df["Start date"].dtype.name == "datetime64[ns]"
+    assert df["Start time"].dtype.name == "timedelta64[ns]"
+    assert df["Duration"].dtype.name == "timedelta64[ns]"
+    assert df["Hours"].dtype.name == "float64"
+
+    df = _prepare_input(toggl_file, column_renames={"Start date": "SD", "Start time": "ST", "Duration": "D"})
+
+    assert "Start date" not in df.columns
+    assert "Start time" not in df.columns
+    assert "Duration" not in df.columns
+    assert df["SD"].dtype.name == "datetime64[ns]"
+    assert df["ST"].dtype.name == "timedelta64[ns]"
+    assert df["D"].dtype.name == "timedelta64[ns]"
+
+
+def test_convert_to_harvest_mocked(toggl_file, spy_files, mock_google_user_info):
+    mock_google_user_info.return_value = {}
+
+    convert_to_harvest(toggl_file, client_name=None)
+
     spy_files.write_csv.assert_called_once()
     call_args = spy_files.write_csv.call_args
     assert sys.stdout in call_args[0]