Fix typing/docstring on dag helper methods (apache#45672)

airflow/models/dag.py

@@ -779,13 +779,13 @@ def get_is_paused(self, session=NEW_SESSION) -> None:
         return session.scalar(select(DagModel.is_paused).where(DagModel.dag_id == self.dag_id))
 
     @provide_session
-    def get_bundle_name(self, session=NEW_SESSION) -> None:
+    def get_bundle_name(self, session=NEW_SESSION) -> str | None:
         """Return the bundle name this DAG is in."""
         return session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == self.dag_id))
 
     @provide_session
-    def get_latest_bundle_version(self, session=NEW_SESSION) -> None:
-        """Return the bundle name this DAG is in."""
+    def get_latest_bundle_version(self, session=NEW_SESSION) -> str | None:
+        """Return the latest version of the bundle this DAG is in."""
         return session.scalar(select(DagModel.latest_bundle_version).where(DagModel.dag_id == self.dag_id))
 
     @methodtools.lru_cache(maxsize=None)

tests/models/test_dag.py

@@ -2089,6 +2089,27 @@ def test_dag_owner_links(self):
         orm_dag_owners = session.query(DagOwnerAttributes).all()
         assert not orm_dag_owners
 
+    def test_get_bundle_name(self, testing_dag_bundle):
+        dag = DAG("dag")
+
+        # until we've sycned, it'll be None
+        assert dag.get_bundle_name() is None
+
+        DAG.bulk_write_to_db("testing", None, [dag])
+        assert dag.get_bundle_name() == "testing"
+
+    def test_get_latest_bundle_version(self, testing_dag_bundle):
+        dag = DAG("dag")
+
+        # until we've sycned, it'll be None
+        assert dag.get_latest_bundle_version() is None
+
+        # Now, it can be none or a str
+        DAG.bulk_write_to_db("testing", None, [dag])
+        assert dag.get_latest_bundle_version() is None
+        DAG.bulk_write_to_db("testing", "abc", [dag])
+        assert dag.get_latest_bundle_version() == "abc"
+
 
 class TestDagModel:
     def _clean(self):