Tests

dags/clean/clean.py

@@ -1,45 +1,46 @@
-from airflow.operators.bash_operator import BashOperator
-from airflow import DAG
+import os
+
 import pendulum
+from airflow import DAG
 from airflow.decorators import dag
 from airflow.operators.bash import BashOperator
-import os
+from airflow.operators.bash_operator import BashOperator
 
 AIRFLOW_HOME = os.getenv("AIRFLOW_HOME")
 LOG_DIR = "logs"
-logs_dir = f'{AIRFLOW_HOME}/{LOG_DIR}'
+logs_dir = f"{AIRFLOW_HOME}/{LOG_DIR}"
 
 default_args = {
-    'owner': 'airflow',
-    'start_date': pendulum.today("UTC").add(days=-1),
-    'schedule': "@monthly",
-    'catchup': False,
-    'max_active_runs': 1,
+    "owner": "airflow",
+    "start_date": pendulum.today("UTC").add(days=-1),
+    "schedule": "@monthly",
+    "catchup": False,
+    "max_active_runs": 1,
 }
 
 dag = DAG(
-    'cleanup_logs',
+    "cleanup_logs",
     default_args=default_args,
-    description='DAG to clean up old log directories',
-    schedule_interval='@monthly',
+    description="DAG to clean up old log directories",
+    schedule_interval="@monthly",
 )
 
 cleanup_command = f"""
     logs_dir="{logs_dir}"
     current_date=$(date "+%Y-%m-%dT%H:%M:%S.%N%z")
 
-    find "$logs_dir" -type f -path "*/run_id=*/task_id=*" -name "attempt=*.log" -exec sh -c '
+    find "$logs_dir" -type d -path "*/run_id=*" -exec sh -c '
         run_id=$(dirname "$1" | grep -oP "run_id=\K[^/]+")
         task_id=$(dirname "$1" | grep -oP "task_id=\K[^/]+")
         if [ $(date -d "$run_id" "+%s") -lt $(date -d "$2" "+%s" --date="-30 days") ]; then
             echo "Deleting: $1"
-            rm "$1"
+            rm -r "$1"
         fi
     ' _ {{}} "$current_date" \;
 """
 
 cleanup_task = BashOperator(
-    task_id='cleanup_logs',
+    task_id="cleanup_logs",
     bash_command=cleanup_command,
     dag=dag,
 )

tests/units/clean/test_clean.py

@@ -0,0 +1,39 @@
+
+import os
+import logging
+import subprocess
+import datetime
+from pytest import fixture
+from unittest import mock
+from freezegun import freeze_time
+
+from clean.clean import cleanup_command
+
+
+@fixture
+@freeze_time("2023-10-20")
+def create_log():
+    logs_date = datetime.datetime.now()
+    log_path = f'tests/units/clean/logs/dag_id=test/run_id=test__{logs_date.strftime("%Y-%m-%dT%H:%M:%S.%f%z")}/task_id=test_task/attempt=1.log'
+    os.makedirs(os.path.dirname(log_path), exist_ok=True)
+    try:
+        with open(log_path, "w") as f:
+            f.write("test log")
+        return True
+    except:
+        logging.error('Failed file creation for tests')
+        pass
+
+@freeze_time("2023-10-20")
+@mock.patch.dict(os.environ, {"LOG_DIR": "tests/units/clean/logs"})
+def test_clean_up_command(create_log):
+    if create_log:
+        assert os.listdir(f'./tests/units/clean/logs/dag_id=test/run_id=test__2023-10-20T00:00:00.000000/task_id=test_task') == ["attempt=1.log"]
+        AIRFLOW_HOME = os.getenv("AIRFLOW_HOME")
+        LOG_DIR = os.getenv("LOG_DIR")
+        logs_dir = f'{AIRFLOW_HOME}/{LOG_DIR}'
+        subprocess.run(cleanup_command, shell=True)
+        for (_, _, filenames) in os.walk(logs_dir):
+            while filenames == 'attempt=1.log':
+                logging.error("Log was not deleted!")
+                assert None