amoghrajesh commented on code in PR #62010:
URL: https://github.com/apache/airflow/pull/62010#discussion_r2812002809
##########
task-sdk/tests/task_sdk/test_lineage.py:
##########
@@ -565,6 +606,34 @@ def test_extra_limit(self, collector):
assert len(collector.collected_assets.extra) == 200
+ def test_configurable_max_collected_extra(self):
+ """MAX_COLLECTED_EXTRA is read from [lineage] conf and limits extra
collection."""
+ import importlib
+
+ import airflow.sdk.lineage as lineage_mod
+ from airflow.sdk.configuration import conf
+
+ real_getint = conf.getint
+ new_max = 4
+
+ with patch.object(
+ conf,
+ "getint",
+ side_effect=lambda s, k, fallback=None: (
+ new_max if k == "hll_max_collected_extra" else real_getint(s,
k, fallback=fallback)
+ ),
+ ):
+ importlib.reload(lineage_mod)
+ assert new_max == lineage_mod.MAX_COLLECTED_EXTRA
+
+ collector = lineage_mod.HookLineageCollector()
+ ctx = MagicMock(spec=BaseHook)
+ for i in range(new_max * 2):
+ collector.add_extra(ctx, f"key_{i}", f"value_{i}")
+ assert len(collector.collected_assets.extra) == new_max
+
+ importlib.reload(lineage_mod)
Review Comment:
Do you need to patch `conf.getint`? Can't you use `@conf_vars` instead and
reload the module like yuo are doing for it to catch the new values?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
