Skip to content

Commit 88e9829

Browse files
cka-ycka-y
committed
test: tdg redirect
1 parent 459898d commit 88e9829

File tree

3 files changed

+260
-6
lines changed

3 files changed

+260
-6
lines changed

‎functions-python/tasks_executor/src/main.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@
2323
from shared.helpers.logger import init_logger
2424
from tasks.data_import.transitfeeds.sync_transitfeeds import sync_transitfeeds_handler
2525
from tasks.data_import.transportdatagouv.import_tdg_feeds import import_tdg_handler
26-
from tasks.data_import.transportdatagouv.redirect_mdb_feeds import update_tdg_redirects_handler
26+
from tasks.data_import.transportdatagouv.update_tdg_redirects import (
27+
update_tdg_redirects_handler,
28+
)
2729
from tasks.dataset_files.rebuild_missing_dataset_files import (
2830
rebuild_missing_dataset_files_handler,
2931
)
@@ -118,7 +120,7 @@
118120
"mdb_to_tdg_redirect": {
119121
"description": "Redirect duplicate MDB feeds to TDG imported feeds.",
120122
"handler": update_tdg_redirects_handler,
121-
}
123+
},
122124
}
123125

124126

functions-python/tasks_executor/src/tasks/data_import/transportdatagouv/redirect_mdb_feeds.py renamed to functions-python/tasks_executor/src/tasks/data_import/transportdatagouv/update_tdg_redirects.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,8 @@ def _update_feed_redirect(
8282
redirect = (
8383
db_session.query(Redirectingid)
8484
.filter(
85-
Redirectingid.target_id == tdg_stable_id,
86-
Redirectingid.source_id == mdb_stable_id,
85+
Redirectingid.target_id == tdg_feed.id,
86+
Redirectingid.source_id == mdb_feed.id,
8787
)
8888
.one_or_none()
8989
)
@@ -103,8 +103,8 @@ def _update_feed_redirect(
103103
tdg_stable_id,
104104
)
105105
redirect = Redirectingid(
106-
target_id=tdg_stable_id,
107-
source_id=mdb_stable_id,
106+
target_id=tdg_feed.id,
107+
source_id=mdb_feed.id,
108108
redirect_comment="Redirecting post TDG import",
109109
)
110110
db_session.add(redirect)

‎functions-python/tasks_executor/tests/tasks/data_import/test_redirect_mdb_feeds.py‎

Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
import os
2+
import unittest
3+
from typing import Optional
4+
from unittest.mock import patch
5+
6+
import pandas as pd
7+
from sqlalchemy.orm import Session
8+
9+
from shared.database.database import with_db_session
10+
from shared.database_gen.sqlacodegen_models import (
11+
Feed,
12+
Gtfsfeed,
13+
Redirectingid,
14+
)
15+
from tasks.data_import.data_import_utils import get_or_create_feed
16+
from tasks.data_import.transportdatagouv.update_tdg_redirects import (
17+
update_tdg_redirects_handler,
18+
_update_feed_redirect,
19+
)
20+
from test_shared.test_utils.database_utils import default_db_url
21+
22+
TEST_STABLE_IDS = {
23+
"mdb-foo",
24+
"tdg-bar",
25+
"mdb-missing",
26+
"tdg-missing",
27+
}
28+
29+
30+
@with_db_session(db_url=default_db_url)
31+
def _cleanup_tdg_redirect_test_data(db_session: Session):
32+
"""
33+
Remove only the feeds / redirects created by this test suite,
34+
identified by their stable_ids.
35+
"""
36+
# Find feeds created for this test suite
37+
feeds = db_session.query(Feed).filter(Feed.stable_id.in_(TEST_STABLE_IDS)).all()
38+
39+
if not feeds:
40+
return
41+
42+
feed_ids = [f.id for f in feeds]
43+
44+
# Delete Redirectingid rows involving these feeds
45+
(
46+
db_session.query(Redirectingid)
47+
.filter(
48+
(Redirectingid.source_id.in_(feed_ids))
49+
| (Redirectingid.target_id.in_(feed_ids))
50+
)
51+
.delete(synchronize_session=False)
52+
)
53+
54+
# Delete the feeds themselves (should cascade to Gtfsfeed, etc., if configured)
55+
for feed in feeds:
56+
db_session.delete(feed)
57+
58+
59+
class TestUpdateTDGRedirectsHelpers(unittest.TestCase):
60+
def tearDown(self):
61+
_cleanup_tdg_redirect_test_data()
62+
63+
@with_db_session(db_url=default_db_url)
64+
def test_update_feed_redirect_creates_redirect(self, db_session: Session):
65+
"""
66+
If both MDB and TDG feeds exist and no redirect is present,
67+
_update_feed_redirect should create one and return proper counters.
68+
"""
69+
mdb_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "mdb-foo", "gtfs")
70+
tdg_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "tdg-bar", "gtfs")
71+
db_session.flush()
72+
73+
counters = _update_feed_redirect(
74+
db_session=db_session,
75+
mdb_stable_id="mdb-foo",
76+
tdg_stable_id="tdg-bar",
77+
)
78+
79+
self.assertEqual(counters["redirects_created"], 1)
80+
self.assertEqual(counters["redirects_existing"], 0)
81+
self.assertEqual(counters["missing_mdb_feeds"], 0)
82+
self.assertEqual(counters["missing_tdg_feeds"], 0)
83+
db_session.commit()
84+
85+
redirect: Optional[Redirectingid] = (
86+
db_session.query(Redirectingid)
87+
.filter(
88+
Redirectingid.source_id == mdb_feed.id,
89+
Redirectingid.target_id == tdg_feed.id,
90+
)
91+
.one_or_none()
92+
)
93+
self.assertIsNotNone(redirect)
94+
self.assertEqual(redirect.redirect_comment, "Redirecting post TDG import")
95+
96+
@with_db_session(db_url=default_db_url)
97+
def test_update_feed_redirect_missing_mdb(self, db_session: Session):
98+
"""
99+
If MDB feed is missing, it should be counted and no redirect created.
100+
"""
101+
tdg_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "tdg-bar", "gtfs")
102+
db_session.flush()
103+
104+
counters = _update_feed_redirect(
105+
db_session=db_session,
106+
mdb_stable_id="mdb-missing",
107+
tdg_stable_id="tdg-bar",
108+
)
109+
110+
self.assertEqual(counters["redirects_created"], 0)
111+
self.assertEqual(counters["redirects_existing"], 0)
112+
self.assertEqual(counters["missing_mdb_feeds"], 1)
113+
self.assertEqual(counters["missing_tdg_feeds"], 0)
114+
115+
redirects = db_session.query(Redirectingid).all()
116+
self.assertEqual(len(redirects), 0)
117+
118+
@with_db_session(db_url=default_db_url)
119+
def test_update_feed_redirect_missing_tdg(self, db_session: Session):
120+
"""
121+
If TDG feed is missing, it should be counted and no redirect created.
122+
"""
123+
mdb_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "mdb-foo", "gtfs")
124+
db_session.flush()
125+
126+
counters = _update_feed_redirect(
127+
db_session=db_session,
128+
mdb_stable_id="mdb-foo",
129+
tdg_stable_id="tdg-missing",
130+
)
131+
132+
self.assertEqual(counters["redirects_created"], 0)
133+
self.assertEqual(counters["redirects_existing"], 0)
134+
self.assertEqual(counters["missing_mdb_feeds"], 0)
135+
self.assertEqual(counters["missing_tdg_feeds"], 1)
136+
137+
redirects = db_session.query(Redirectingid).all()
138+
self.assertEqual(len(redirects), 0)
139+
140+
@with_db_session(db_url=default_db_url)
141+
def test_update_feed_redirect_existing_redirect(self, db_session: Session):
142+
"""
143+
If redirect already exists, it should be counted as existing
144+
and no new row created.
145+
"""
146+
mdb_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "mdb-foo", "gtfs")
147+
tdg_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "tdg-bar", "gtfs")
148+
db_session.flush()
149+
150+
existing = Redirectingid(
151+
source_id=mdb_feed.id,
152+
target_id=tdg_feed.id,
153+
redirect_comment="Existing redirect",
154+
)
155+
db_session.add(existing)
156+
db_session.flush()
157+
158+
counters = _update_feed_redirect(
159+
db_session=db_session,
160+
mdb_stable_id="mdb-foo",
161+
tdg_stable_id="tdg-bar",
162+
)
163+
164+
self.assertEqual(counters["redirects_created"], 0)
165+
self.assertEqual(counters["redirects_existing"], 1)
166+
self.assertEqual(counters["missing_mdb_feeds"], 0)
167+
self.assertEqual(counters["missing_tdg_feeds"], 0)
168+
169+
redirects = db_session.query(Redirectingid).all()
170+
self.assertEqual(len(redirects), 1)
171+
172+
173+
class TestUpdateTDGRedirectsHandler(unittest.TestCase):
174+
def tearDown(self):
175+
_cleanup_tdg_redirect_test_data()
176+
177+
@with_db_session(db_url=default_db_url)
178+
def test_handler_creates_redirects_from_csv(self, db_session: Session):
179+
"""
180+
Happy path:
181+
- CSV has one row with MDB ID + raw TDG ID
182+
- MDB and corresponding TDG feeds exist
183+
- handler creates one redirect and commits it
184+
"""
185+
mdb_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "mdb-foo", "gtfs")
186+
tdg_feed, _ = get_or_create_feed(db_session, Gtfsfeed, "tdg-bar", "gtfs")
187+
db_session.commit()
188+
189+
df = pd.DataFrame(
190+
[
191+
{
192+
"MDB ID": "mdb-foo",
193+
"TDG ID": "bar", # maps to "tdg-bar"
194+
}
195+
]
196+
)
197+
198+
with patch(
199+
"tasks.data_import.transportdatagouv.update_tdg_redirects.pd.read_csv",
200+
return_value=df,
201+
), patch.dict(
202+
os.environ,
203+
{"COMMIT_BATCH_SIZE": "1"},
204+
clear=False,
205+
):
206+
result = update_tdg_redirects_handler({"dry_run": False})
207+
208+
self.assertEqual(
209+
result,
210+
{
211+
"message": "TDG redirects update executed successfully.",
212+
"rows_processed": 1,
213+
"redirects_created": 1,
214+
"redirects_existing": 0,
215+
"missing_mdb_feeds": 0,
216+
"missing_tdg_feeds": 0,
217+
"params": {"dry_run": False},
218+
},
219+
)
220+
221+
redirect: Optional[Redirectingid] = (
222+
db_session.query(Redirectingid)
223+
.filter(
224+
Redirectingid.source_id == mdb_feed.id,
225+
Redirectingid.target_id == tdg_feed.id,
226+
)
227+
.one_or_none()
228+
)
229+
self.assertIsNotNone(redirect)
230+
231+
def test_handler_csv_load_failure(self):
232+
"""
233+
If pd.read_csv throws, handler should return an error summary
234+
and not raise.
235+
"""
236+
with patch(
237+
"tasks.data_import.transportdatagouv.update_tdg_redirects.pd.read_csv",
238+
side_effect=RuntimeError("boom"),
239+
):
240+
result = update_tdg_redirects_handler({"dry_run": True})
241+
242+
self.assertEqual(result["message"], "Failed to load TDG redirect CSV.")
243+
self.assertIn("error", result)
244+
self.assertEqual(result["rows_processed"], 0)
245+
self.assertEqual(result["redirects_created"], 0)
246+
self.assertEqual(result["redirects_existing"], 0)
247+
self.assertEqual(result["missing_mdb_feeds"], 0)
248+
self.assertEqual(result["missing_tdg_feeds"], 0)
249+
250+
251+
if __name__ == "__main__":
252+
unittest.main()

0 commit comments

Comments
 (0)