test: temp changes to check_connections DAG for stress test

Author: helen-m-lin

src/aind_airflow_jobs/dag_tasks/check_connections.py

@@ -5,6 +5,7 @@
 import os
 import sys
 from pathlib import Path
+from time import sleep
 
 import boto3
 
@@ -16,6 +17,8 @@
     stream=sys.stdout,
 )
 
+SLEEP_TIME = 120
+
 
 class CheckConnectionsDag(DagTasks):
     """DAG tasks for checking connections to external services."""
@@ -43,6 +46,7 @@ def check_param_store_connection(self):
             raise AssertionError("Unable to retrieve ams_uri!")
         if not co_uri:
             raise AssertionError("Unable to retrieve co_uri!")
+        sleep(SLEEP_TIME)
 
     def check_aws_connection(self):
         """Check AWS S3 connection."""
@@ -59,6 +63,7 @@ def check_aws_connection(self):
             )
         finally:
             s3_client.close()
+        sleep(SLEEP_TIME)
 
     def check_slurm_connection(self):
         """Check SLURM connection."""
@@ -81,6 +86,7 @@ def check_vast_connection(self):
         is_dir = Path(mounted_directory).is_dir()
         if not is_dir:
             raise NotADirectoryError(f"{mounted_directory} not recognized!")
+        sleep(SLEEP_TIME)
 
     def check_hpc_connection(self):
         """Check hpc ssh command output can be read"""

tests/dag_tasks/test_check_connections.py

@@ -22,6 +22,7 @@ def setUpClass(cls):
         dag = CheckConnectionsDag(airflow_task_settings=settings)
         cls.dag = dag
 
+    @patch("aind_airflow_jobs.dag_tasks.check_connections.sleep")
     @patch.dict(
         "os.environ",
         {
@@ -31,7 +32,7 @@ def setUpClass(cls):
         },
         clear=True,
     )
-    def test_check_param_store_connection(self):
+    def test_check_param_store_connection(self, mock_sleep: MagicMock):
         """Tests check_param_store_connection."""
 
         with self.assertLogs(level="INFO") as captured:
@@ -44,6 +45,7 @@ def test_check_param_store_connection(self):
             ],
             captured.output,
         )
+        mock_sleep.assert_called_once()
 
     @patch.dict(
         "os.environ",
@@ -122,8 +124,11 @@ def test_check_param_store_connection_missing_co_uri(self):
 
         self.assertEqual("Unable to retrieve co_uri!", exc.exception.args[0])
 
+    @patch("aind_airflow_jobs.dag_tasks.check_connections.sleep")
     @patch("aind_airflow_jobs.dag_tasks.check_connections.boto3.client")
-    def test_check_aws_connection(self, mock_boto_client: MagicMock):
+    def test_check_aws_connection(
+        self, mock_boto_client: MagicMock, mock_sleep: MagicMock
+    ):
         """Tests check_aws_connection."""
 
         mock_s3_client = MagicMock()
@@ -136,6 +141,7 @@ def test_check_aws_connection(self, mock_boto_client: MagicMock):
             Bucket="my-bucket", MaxKeys=1
         )
         mock_s3_client.close.assert_called_once()
+        mock_sleep.assert_called_once()
 
     @patch("aind_airflow_jobs.dag_tasks.check_connections.boto3.client")
     def test_check_aws_connection_closes_on_error(
@@ -189,14 +195,18 @@ def test_check_slurm_connection(
             captured.output,
         )
 
+    @patch("aind_airflow_jobs.dag_tasks.check_connections.sleep")
     @patch("aind_airflow_jobs.dag_tasks.check_connections.Path.is_dir")
-    def test_check_vast_connection(self, mock_is_dir: MagicMock):
+    def test_check_vast_connection(
+        self, mock_is_dir: MagicMock, mock_sleep: MagicMock
+    ):
         """Tests check_vast_connection validates mounted directory."""
 
         mock_is_dir.return_value = True
         self.dag.check_vast_connection()
 
         mock_is_dir.assert_called_once_with()
+        mock_sleep.assert_called_once()
 
     @patch("aind_airflow_jobs.dag_tasks.check_connections.Path.is_dir")
     def test_check_vast_connection_not_directory(self, mock_is_dir: MagicMock):