vincbeck commented on code in PR #41004:
URL: https://github.com/apache/airflow/pull/41004#discussion_r1690214528
##########
airflow/settings.py:
##########
@@ -281,6 +281,59 @@ def remove(*args, **kwargs):
pass
+AIRFLOW_PATH_FOR_TESTS = os.path.dirname(os.path.dirname(__file__))
+AIRFLOW_TESTS_PATH_FOR_TESTS = os.path.join(AIRFLOW_PATH_FOR_TESTS, "tests")
+
+
+class TracebackSessionForTests:
+ """
+ Session that throws error when you try to create a session outside of the
test code.
+
+ When we run our tests in "db isolation" mode we except that "airflow" code
will never create
+ a session on its own and internal_api server is used for all calls but the
test code might use
+ the session to setup and teardown in the DB so that the internal API
server accesses it.
+
+ :meta private:
+ """
+
+ db_session_class = None
+
+ def __init__(self):
+ self.traceback = traceback.extract_stack()
+ self.current_db_session = TracebackSessionForTests.db_session_class()
+
+ def __getattr__(self, item):
+ if self.is_called_from_test_code():
+ return getattr(self.current_db_session, item)
+ raise RuntimeError(
+ "TracebackSessionForTests object was used but internal API is
enabled. "
+ "Only test code is allowed to use this object. "
+ "You'll need to ensure you are making only RPC calls with this
object. "
+ "The stack list below will show where the TracebackSession object
was created."
+ + "\n".join(traceback.format_list(self.traceback))
+ )
+
+ def remove(*args, **kwargs):
+ pass
+
+ def is_called_from_test_code(self) -> bool:
+ """
+ Check if the object was created from test code.
+
+ This is done by checking if the first "airflow" filename in the
traceback
+ is "airflow/tests" or "regular airflow".
+
+ :meta: private
+ :return: True if the object was created from test code, False
otherwise.
+ """
+ for tb in self.traceback:
+ if tb.filename.startswith(AIRFLOW_PATH_FOR_TESTS):
+ if tb.filename.startswith(AIRFLOW_TESTS_PATH_FOR_TESTS):
+ return True
+ return False
Review Comment:
```suggestion
return tb.filename.startswith(AIRFLOW_TESTS_PATH_FOR_TESTS):
```
--
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]
