From 1a535ee918f6bed2e78d0d36b443bf9463f33f47 Mon Sep 17 00:00:00 2001
From: Tim Paine <3105306+timkpaine@users.noreply.github.com>
Date: Sat, 15 Feb 2025 23:27:28 -0500
Subject: [PATCH] Starting on test coverage

---
 README.md                            |  6 +++---
 airflow_supervisor/airflow/local.py  |  6 +++---
 airflow_supervisor/tests/test_dag.py | 24 ++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)
 create mode 100644 airflow_supervisor/tests/test_dag.py

diff --git a/README.md b/README.md
index 1e2f027..5558fe8 100755
--- a/README.md
+++ b/README.md
@@ -2,9 +2,9 @@
 
 [Apache Airflow](https://airflow.apache.org) utilities for running long-running or always-on jobs with [supervisord](http://supervisord.org)
 
-[![Build Status](https://github.com/timkpaine/airflow-supervisor/actions/workflows/build.yml/badge.svg)](https://github.com/timkpaine/airflow-supervisor/actions?query=workflow%3A%22Build+Status%22)
-[![codecov](https://codecov.io/gh/timkpaine/airflow-supervisor/branch/main/graph/badge.svg)](https://codecov.io/gh/timkpaine/airflow-supervisor)
-[![License](https://img.shields.io/github/license/timkpaine/airflow-supervisor)](https://github.com/timkpaine/airflow-supervisor)
+[![Build Status](https://github.com/airflow-laminar/airflow-supervisor/actions/workflows/build.yml/badge.svg)](https://github.com/airflow-laminar/airflow-supervisor/actions?query=workflow%3A%22Build+Status%22)
+[![codecov](https://codecov.io/gh/airflow-laminar/airflow-supervisor/branch/main/graph/badge.svg)](https://codecov.io/gh/airflow-laminar/airflow-supervisor)
+[![License](https://img.shields.io/github/license/airflow-laminar/airflow-supervisor)](https://github.com/airflow-laminar/airflow-supervisor)
 [![PyPI](https://img.shields.io/pypi/v/airflow-supervisor.svg)](https://pypi.python.org/pypi/airflow-supervisor)
 
 ## Overview
diff --git a/airflow_supervisor/airflow/local.py b/airflow_supervisor/airflow/local.py
index ff28a60..b590e6a 100644
--- a/airflow_supervisor/airflow/local.py
+++ b/airflow_supervisor/airflow/local.py
@@ -25,12 +25,12 @@
 
 
 class Supervisor(object):
-    _dag: "DAG"
+    _dag: DAG
     _cfg: SupervisorAirflowConfiguration
-    _kill_dag: "DAG"
+    _kill_dag: DAG
     _xmlrpc_client: SupervisorRemoteXMLRPCClient
 
-    def __init__(self, dag: "DAG", cfg: SupervisorAirflowConfiguration, **kwargs):
+    def __init__(self, dag: DAG, cfg: SupervisorAirflowConfiguration, **kwargs):
         # store config
         self._cfg = cfg
 
diff --git a/airflow_supervisor/tests/test_dag.py b/airflow_supervisor/tests/test_dag.py
new file mode 100644
index 0000000..7f9b71f
--- /dev/null
+++ b/airflow_supervisor/tests/test_dag.py
@@ -0,0 +1,24 @@
+from airflow.models.dag import DAG
+
+from airflow_supervisor import Supervisor, SupervisorAirflowConfiguration
+
+
+class TestDag:
+    def test_instantiation(self, supervisor_airflow_configuration: SupervisorAirflowConfiguration):
+        dag = DAG(dag_id="test_dag", default_args={}, schedule=None, params={})
+        s = Supervisor(dag=dag, cfg=supervisor_airflow_configuration)
+        assert len(dag.tasks) == 17
+        assert dag.catchup is False
+        assert dag.concurrency == 1
+        assert dag.max_active_tasks == 1
+        assert dag.max_active_runs == 1
+
+        assert s.configure_supervisor in dag.tasks
+        assert s.configure_supervisor in dag.tasks
+        assert s.start_supervisor in dag.tasks
+        assert s.start_programs in dag.tasks
+        assert s.check_programs in dag.tasks
+        assert s.restart_programs in dag.tasks
+        assert s.stop_programs in dag.tasks
+        assert s.stop_supervisor in dag.tasks
+        assert s.unconfigure_supervisor in dag.tasks