"""Unit tests for DAG run time logic.""" import datetime as dt import json from unittest.mock import patch import pytest from abacus_contract.logic import dag_run_times _iso = dag_run_times._iso runtimes_from_stdout = dag_run_times.runtimes_from_stdout get_average_dag_run_time = dag_run_times.get_average_dag_run_time def test__iso_parses_z_and_offset_equivalence(): """Ensure Z and +00:00 offsets are treated the same.""" z = '2025-09-01T12:34:56.123Z' off = '2025-09-01T13:34:56.123+01:00' dt_z = _iso(z) dt_off = _iso(off) assert (dt_off - dt_z).total_seconds() == 0 assert dt_z.tzinfo is not None and dt_off.tzinfo is not None def test_runtimes_from_stdout_empty_and_missing_dates(): """No runs, or runs with missing start/end dates.""" runs = [ {'start_date': None, 'end_date': None}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': None}, {'start_date': None, 'end_date': '2025-09-01T00:10:00Z'}, ] out = runtimes_from_stdout(json.dumps(runs)) assert out == { 'count': 0, 'avg_seconds': None, 'median_seconds': None, 'p95_seconds': None, } def test_runtimes_from_stdout_basic_stats_even_count(): """Durations: 60, 120, 180, 240.""" runs = [ {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:01:00Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:02:00Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:03:00Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:04:00Z'}, ] out = runtimes_from_stdout(json.dumps(runs)) assert out['count'] == 4 assert out['min_run_time_seconds'] == 60 assert out['max_run_time_seconds'] == 240 assert out['average_run_time_seconds'] == pytest.approx(150.0) assert out['median_run_time_seconds'] == pytest.approx(150.0) # p95 index = ceil(0.95*4)-1 = 3 -> last element (240) assert out['p95_run_time_seconds'] == 240 def test_runtimes_from_stdout_basic_stats_odd_count(): """Durations: 10, 20, 30, 40, 50.""" runs = [ {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:10Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:20Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:30Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:40Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:50Z'}, ] out = runtimes_from_stdout(json.dumps(runs)) assert out['count'] == 5 assert out['median_run_time_seconds'] == 30 # p95 index = ceil(0.95*5)-1 = 4 -> last element (50) assert out['p95_run_time_seconds'] == 50 def test_runtimes_from_stdout_mixed_order_and_decimals(): """Durations: 60, 30, 30.5, 120, 60.""" durs = [60, 30, 30.5, 120, 60] runs = [ {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:01:00Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:30Z'}, { 'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:00:30.500Z', }, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:02:00Z'}, {'start_date': '2025-09-01T00:00:00Z', 'end_date': '2025-09-01T00:01:00Z'}, ] out = runtimes_from_stdout(json.dumps(runs)) d_sorted = sorted(durs) assert out['count'] == 5 assert out['min_run_time_seconds'] == d_sorted[0] assert out['max_run_time_seconds'] == d_sorted[-1] assert out['median_run_time_seconds'] == d_sorted[2] assert out['average_run_time_seconds'] == pytest.approx(sum(durs) / 5) assert out['p95_run_time_seconds'] == d_sorted[-1] class _FixedNow(dt.datetime): @classmethod def now(cls, tz=None): # 2025-09-26 10:00:00 (naive; code appends 'Z') return cls(2025, 9, 26, 10, 0, 0, tzinfo=None) def test_get_average_dag_run_time_success(): """Simulate successful path with two runs.""" dag_id = 'accounting_run_calculate' with ( patch( 'abacus_contract.logic.dag_run_times.dt.datetime', _FixedNow, create=True ), patch( 'abacus_contract.logic.dag_run_times.run_airflow_cli_command' ) as mock_cli, patch('abacus_contract.logic.dag_run_times.response.Response') as MockResponse, patch( 'abacus_contract.logic.dag_run_times.response.create_error_response' ) as mock_err_resp, ): runs = [ { 'start_date': '2025-09-25T10:00:00Z', 'end_date': '2025-09-25T10:10:00Z', }, # 600 { 'start_date': '2025-09-24T10:00:00Z', 'end_date': '2025-09-24T10:20:00Z', }, # 1200 ] mock_cli.return_value = {'errors': '', 'output': json.dumps(runs)} captured = {} def _capture(message=None, status=None, **_): captured['message'] = message captured['status'] = status class _R: ... return _R() MockResponse.side_effect = _capture resp = get_average_dag_run_time(dag_id) assert mock_cli.call_count == 1 [cmd] = mock_cli.call_args[0] assert f'dags list-runs -d {dag_id}' in cmd assert '--state success --output json' in cmd assert '-e 2025-09-26T10:00:00Z' in cmd assert '-s 2025-08-27T10:00:00Z' in cmd # 30 days back # Ensure error path NOT used mock_err_resp.assert_not_called() # Payload checks assert captured['status'] == 200 msg = captured['message'] assert msg['dag_id'] == dag_id assert msg['count'] == 2 assert msg['min_run_time_seconds'] == 600 assert msg['max_run_time_seconds'] == 1200 assert msg['median_run_time_seconds'] == pytest.approx((600 + 1200) / 2) assert msg['average_run_time_seconds'] == pytest.approx((600 + 1200) / 2) assert msg['p95_run_time_seconds'] == 1200 assert resp is not None def test_get_average_dag_run_time_error_path(): """Simulate error from Airflow CLI command.""" dag_id = 'missing_dag' with ( patch( 'abacus_contract.logic.dag_run_times.run_airflow_cli_command' ) as mock_cli, patch( 'abacus_contract.logic.dag_run_times.response.create_error_response' ) as mock_err_resp, ): mock_cli.return_value = {'errors': 'DAG not found', 'output': ''} sentinel = object() def _err(**kwargs): assert kwargs['code'] == dag_run_times.error.ERROR_DAG_RUN_TIMES assert kwargs['status'] == 500 assert 'Error getting DAG run times' in kwargs['message'] return sentinel mock_err_resp.side_effect = _err resp = get_average_dag_run_time(dag_id) mock_cli.assert_called_once() mock_err_resp.assert_called_once() assert resp is sentinel