open-metadata--openmetadata/ingestion/tests/unit/metadata/utils/test_progress_tracker.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Tests for the ProgressTrackerState singleton"""

import threading

import pytest

from metadata.utils.progress_tracker import EntityProgress, ProgressTrackerState
from metadata.utils.singleton import Singleton


class TestEntityProgress:
    """Tests for EntityProgress model"""

    def test_estimate_remaining_seconds_returns_none_when_no_times(self) -> None:
        progress = EntityProgress(total=100, processed=10)
        assert progress.estimate_remaining_seconds() is None

    def test_estimate_remaining_seconds_returns_none_when_complete(self) -> None:
        progress = EntityProgress(total=100, processed=100, processing_times=[0.1, 0.2, 0.15])
        assert progress.estimate_remaining_seconds() is None

    def test_estimate_remaining_seconds_calculates_correctly(self) -> None:
        progress = EntityProgress(total=100, processed=50, processing_times=[1.0, 1.0, 1.0])
        remaining = progress.estimate_remaining_seconds()
        assert remaining == 50  # 50 remaining * 1.0 avg time

    def test_estimate_remaining_uses_rolling_window(self) -> None:
        times = [10.0] * 50 + [1.0] * 100
        progress = EntityProgress(total=200, processed=100, processing_times=times)
        remaining = progress.estimate_remaining_seconds()
        assert remaining == 100  # Uses last 100 times (avg 1.0) * 100 remaining

    def test_get_processing_rate_returns_none_when_no_times(self) -> None:
        progress = EntityProgress(total=100, processed=10)
        assert progress.get_processing_rate() is None

    def test_get_processing_rate_calculates_correctly(self) -> None:
        progress = EntityProgress(total=100, processed=50, processing_times=[0.5, 0.5, 0.5])
        rate = progress.get_processing_rate()
        assert rate == 2.0  # 1 / 0.5 = 2 entities per second

    def test_to_dict_format(self) -> None:
        progress = EntityProgress(total=100, processed=25, processing_times=[2.0, 2.0, 2.0])
        result = progress.to_dict()
        assert result["total"] == 100
        assert result["processed"] == 25
        assert result["estimatedRemainingSeconds"] == 150  # 75 remaining * 2.0 avg


class TestProgressTrackerState:
    """Tests for ProgressTrackerState singleton"""

    @pytest.fixture(autouse=True)
    def reset_singleton(self):
        """Reset singleton state before each test"""
        Singleton.clear_all()
        yield
        Singleton.clear_all()

    def test_singleton_returns_same_instance(self) -> None:
        tracker1 = ProgressTrackerState()
        tracker2 = ProgressTrackerState()
        assert tracker1 is tracker2

    def test_set_total_initializes_progress(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)

        progress = tracker.get_progress("Table")
        assert progress is not None
        assert progress.total == 100
        assert progress.processed == 0
        assert progress.start_time is not None

    def test_set_total_resets_existing_progress(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 50)
        tracker.increment_processed("Table", 0.1)

        tracker.set_total("Table", 100)
        progress = tracker.get_progress("Table")
        assert progress.total == 100
        assert progress.processed == 0
        assert len(progress.processing_times) == 0

    def test_add_to_total_increments_total(self) -> None:
        tracker = ProgressTrackerState()
        tracker.add_to_total("Table", 50)
        tracker.add_to_total("Table", 30)

        progress = tracker.get_progress("Table")
        assert progress.total == 80

    def test_increment_processed_updates_count(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)

        tracker.increment_processed("Table", 0.1)
        tracker.increment_processed("Table", 0.2)

        progress = tracker.get_progress("Table")
        assert progress.processed == 2
        assert len(progress.processing_times) == 2

    def test_increment_processed_creates_entry_if_not_exists(self) -> None:
        tracker = ProgressTrackerState()
        tracker.increment_processed("Table", 0.1)

        progress = tracker.get_progress("Table")
        assert progress is not None
        assert progress.processed == 1

    def test_increment_processed_without_time(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)
        tracker.increment_processed("Table")

        progress = tracker.get_progress("Table")
        assert progress.processed == 1
        assert len(progress.processing_times) == 0

    def test_rolling_window_limits_processing_times(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 200)

        for i in range(150):
            tracker.increment_processed("Table", float(i))

        progress = tracker.get_progress("Table")
        assert progress.processed == 150
        assert len(progress.processing_times) == 100

    def test_get_progress_returns_none_for_unknown_entity(self) -> None:
        tracker = ProgressTrackerState()
        assert tracker.get_progress("Unknown") is None

    def test_get_progress_returns_copy(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)

        progress1 = tracker.get_progress("Table")
        progress1.processed = 999  # Modify the copy

        progress2 = tracker.get_progress("Table")
        assert progress2.processed == 0  # Original unchanged

    def test_get_all_progress_returns_all_entities(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)
        tracker.set_total("Database", 10)
        tracker.set_total("Schema", 50)

        all_progress = tracker.get_all_progress()
        assert len(all_progress) == 3
        assert "Table" in all_progress
        assert "Database" in all_progress
        assert "Schema" in all_progress

    def test_get_progress_as_dict_format(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)
        tracker.increment_processed("Table", 1.0)

        result = tracker.get_progress_as_dict()
        assert "Table" in result
        assert result["Table"]["total"] == 100
        assert result["Table"]["processed"] == 1
        assert "estimatedRemainingSeconds" in result["Table"]

    def test_reset_clears_all_progress(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)
        tracker.set_total("Database", 10)

        tracker.reset()

        assert tracker.get_progress("Table") is None
        assert tracker.get_progress("Database") is None
        assert len(tracker.get_all_progress()) == 0

    def test_reset_entity_clears_specific_entity(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 100)
        tracker.set_total("Database", 10)

        tracker.reset_entity("Table")

        assert tracker.get_progress("Table") is None
        assert tracker.get_progress("Database") is not None

    def test_reset_entity_does_not_raise_for_unknown_entity(self) -> None:
        tracker = ProgressTrackerState()
        tracker.reset_entity("Unknown")  # Should not raise

    def test_thread_safety(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Table", 1000)
        errors = []

        def increment_many():
            try:
                for _ in range(100):
                    tracker.increment_processed("Table", 0.001)
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=increment_many) for _ in range(10)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0
        progress = tracker.get_progress("Table")
        assert progress.processed == 1000

    def test_multiple_entity_types(self) -> None:
        tracker = ProgressTrackerState()
        tracker.set_total("Database", 5)
        tracker.set_total("Schema", 20)
        tracker.set_total("Table", 100)

        tracker.increment_processed("Database", 0.5)
        tracker.increment_processed("Schema", 0.2)
        tracker.increment_processed("Table", 0.1)

        assert tracker.get_progress("Database").processed == 1
        assert tracker.get_progress("Schema").processed == 1
        assert tracker.get_progress("Table").processed == 1