Merge pull request #153 from airtai/dev

2 PRs
airtai · Oct 31, 2024 · eb4fed2 · eb4fed2
2 parents f3b7ef9 + 2233b80
commit eb4fed2
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 3 deletions.
diff --git a/google_sheets/data_processing/processing.py b/google_sheets/data_processing/processing.py
@@ -177,6 +177,23 @@ def _replace_values(
     return new_row
 
 
+def _replace_headline_values(new_row: pd.Series, station: Dict[str, Any]) -> pd.Series:
+    # Filter columns that start with "Headline"
+    headline_columns = [col for col in new_row.index if col.startswith("Headline")]
+
+    # Perform replacements only in the headline columns
+    for col in headline_columns:
+        new_row[col] = (
+            new_row[col]
+            .replace(INSERT_STATION_FROM, station["Station From"])
+            .replace(INSERT_STATION_TO, station["Station To"])
+        )
+    return new_row
+
+
+USE_ORIGINAL_STATION_FROM = ["Transfer", "Taxi"]
+
+
 def _process_row(
     new_campaign_row: pd.Series,
     template_row: pd.Series,
@@ -214,6 +231,10 @@ def _process_row(
             language_code=new_row["Language Code"],
             include_locations=include_locations,
         )
+        if new_row["Category"] in USE_ORIGINAL_STATION_FROM:
+            # Use the original Station From in headlines for both directions
+            new_row = _replace_headline_values(new_row, stations[0])
+
         new_row = _replace_values(new_campaign_row, new_row, station)
 
         if target_resource == "ad":

diff --git a/tests/data_processing/test_processing.py b/tests/data_processing/test_processing.py
@@ -7,6 +7,7 @@
     _copy_all_with_prefixes,
     _get_target_location,
     _process_row,
+    _replace_headline_values,
     _replace_values,
     _update_campaign_name,
     _validate_language_codes,
@@ -58,13 +59,39 @@ def test_validate_input_data(df: pd.DataFrame, expected: str) -> None:
     assert validate_input_data(df, mandatory_columns, "name") == expected
 
 
+def test_replace_headline_values() -> None:
+    row = pd.Series(
+        {
+            "Name": "{INSERT_STATION_FROM}",
+            "Headline 1": "{INSERT_STATION_FROM} - {INSERT_STATION_TO}",
+            "Headline 2": "{INSERT_STATION_FROM}",
+        }
+    )
+    station = {
+        "Station From": "A",
+        "Station To": "B",
+    }
+
+    row = _replace_headline_values(row, station)
+    expected_row = pd.Series(
+        {
+            "Name": "{INSERT_STATION_FROM}",
+            "Headline 1": "A - B",
+            "Headline 2": "A",
+        }
+    )
+
+    assert row.equals(expected_row)
+
+
 @pytest.mark.parametrize(
     ("category", "expected_length"),
     [
         (
             "Bus",
             1,
         ),
+        ("Transfer", 1),
     ],
 )
 def test_process_row(
@@ -82,10 +109,10 @@ def test_process_row(
             "Level": "",
             "Keyword Match Type": "Exact",
             "Match Type": "Exact",
-            "Category": "Bus",
+            "Category": category,
             "Target Category": "False",
-            "Ad Group Category": "Bus",
-            "Real Category": "Bus",
+            "Ad Group Category": category,
+            "Real Category": category,
         }
     )
     new_campaign_row = pd.Series(