cleverswarm-python-client/tests/test_cswarm_benchmark_client_extended.py

"""
Extended unit tests for cswarm_benchmark_client module to achieve higher coverage.

Copyright (c) 2016 - present Syncleus, Inc.
Copyright (c) 2016 - present CleverThis, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

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.
"""

import tempfile
from pathlib import Path
from unittest.mock import Mock, patch
import pytest

from cleverswarm_python_client.cswarm_benchmark_client import BenchmarkCLI, main
from cleverswarm_python_client.libs.exceptions import ClientException


class TestBenchmarkCLIExtended:
    """Extended test cases for BenchmarkCLI to achieve higher coverage."""

    def setup_method(self):
        """Set up test CLI."""
        with tempfile.TemporaryDirectory() as temp_dir:
            self.cli = BenchmarkCLI(temp_dir)

    def test_evaluate_results_locally_with_valid_ontologies(self):
        """Test local evaluation with valid ontology IDs."""
        with patch.object(self.cli, 'find_invalid_ontologies') as mock_find:
            mock_find.return_value = []
            with patch('cleverswarm_python_client.cswarm_benchmark_client.EvaluateResults') as mock_eval:
                mock_eval.return_value = {"results": "test"}
                result = self.cli.evaluate_results_locally([1, 2])
                assert result is not None

    def test_print_server_jobs_with_empty_list(self):
        """Test printing server jobs with empty list."""
        with patch.object(self.cli, 'list_server_jobs') as mock_list:
            mock_list.return_value = []
            self.cli.print_server_jobs()
            # Should not raise any exception

    def test_print_server_jobs_detailed_with_wildcards(self):
        """Test printing server jobs with detailed output including wildcards."""
        mock_jobs = [
            {
                "id": "job1",
                "created": "2023-01-01T00:00:00Z",
                "type": "UnstructuredWithWildcards",
                "status": "Completed",
                "retries_count": 0,
                "file_type": "JSONL",
                "machine_id": "machine-001",
                "worker_id": "worker-001",
                "unstructured_sources": ["test.txt"],
                "ontology_sources": ["test.json"],
                "ontology_spec_sources": ["test.owl"],
                "wildcards_sources": ["test.wildcards"],
                "ground_truth_sources": ["test_gt.jsonl"]
            }
        ]

        with patch.object(self.cli, 'list_server_jobs') as mock_list:
            mock_list.return_value = mock_jobs
            self.cli.print_server_jobs(detailed=True)

    def test_print_server_jobs_detailed_without_wildcards(self):
        """Test printing server jobs with detailed output without wildcards."""
        mock_jobs = [
            {
                "id": "job1",
                "created": "2023-01-01T00:00:00Z",
                "type": "UnstructuredWithOntology",
                "status": "Completed",
                "retries_count": 0,
                "file_type": "JSONL",
                "machine_id": "machine-001",
                "worker_id": "worker-001",
                "unstructured_sources": ["test.txt"],
                "ontology_sources": ["test.json"],
                "ontology_spec_sources": ["test.owl"],
                "ground_truth_sources": ["test_gt.jsonl"]
            }
        ]

        with patch.object(self.cli, 'list_server_jobs') as mock_list:
            mock_list.return_value = mock_jobs
            self.cli.print_server_jobs(detailed=True)

    def test_create_benchmark_with_single_ontology(self):
        """Test creating benchmark with single ontology."""
        with patch.object(self.cli, '_client') as mock_client:
            mock_client.create_benchmark_job.return_value = "test_job_123"
            mock_client.append_to_benchmark_job.return_value = True
            
            with tempfile.TemporaryDirectory() as temp_dir:
                self.cli._base_path = Path(temp_dir)
                
                # Create test files
                test_file = Path(temp_dir) / "input_data" / "test_new" / "ont_university_test.jsonl"
                test_file.parent.mkdir(parents=True, exist_ok=True)
                test_file.write_text('{"test": "data"}')
                
                onto_file = Path(temp_dir) / "input_data" / "ontologies_enriched" / "university_ontology_with_descriptions.json"
                onto_file.parent.mkdir(parents=True, exist_ok=True)
                onto_file.write_text('{"ontology": "data"}')
                
                gt_file = Path(temp_dir) / "input_data" / "ground_truth_new" / "ont_university_ground_truth.jsonl"
                gt_file.parent.mkdir(parents=True, exist_ok=True)
                gt_file.write_text('{"ground_truth": "data"}')
                
                result = self.cli.create_benchmark([1])
                assert result == "test_job_123"

    def test_create_benchmark_with_multiple_ontologies(self):
        """Test creating benchmark with multiple ontologies."""
        with patch.object(self.cli, '_client') as mock_client:
            mock_client.create_benchmark_job.return_value = "test_job_123"
            mock_client.append_to_benchmark_job.return_value = True
            
            with tempfile.TemporaryDirectory() as temp_dir:
                self.cli._base_path = Path(temp_dir)
                
                # Create test files for multiple ontologies
                for onto_id in [1, 2]:
                    onto_name = self.cli._all_ontologies[onto_id]
                    
                    test_file = Path(temp_dir) / "input_data" / "test_new" / f"ont_{onto_name}_test.jsonl"
                    test_file.parent.mkdir(parents=True, exist_ok=True)
                    test_file.write_text('{"test": "data"}')
                    
                    onto_file = Path(temp_dir) / "input_data" / "ontologies_enriched" / f"{onto_name}_ontology_with_descriptions.json"
                    onto_file.parent.mkdir(parents=True, exist_ok=True)
                    onto_file.write_text('{"ontology": "data"}')
                    
                    gt_file = Path(temp_dir) / "input_data" / "ground_truth_new" / f"ont_{onto_name}_ground_truth.jsonl"
                    gt_file.parent.mkdir(parents=True, exist_ok=True)
                    gt_file.write_text('{"ground_truth": "data"}')
                
                result = self.cli.create_benchmark([1, 2])
                assert result == "test_job_123"

    def test_download_job_results_with_multiple_ontologies(self):
        """Test downloading job results with multiple ontologies."""
        with patch.object(self.cli, '_client') as mock_client:
            mock_client.retrieve_benchmark_files.return_value = True
            
            with tempfile.TemporaryDirectory() as temp_dir:
                self.cli._base_path = Path(temp_dir)
                
                self.cli.download_job_results("test_job_123", [1, 2])
                
                # Verify that retrieve_benchmark_files was called
                mock_client.retrieve_benchmark_files.assert_called_once()



class TestBenchmarkCLIMainFunctionExtended:
    """Extended test cases for BenchmarkCLI main function to achieve higher coverage."""

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_delete_and_exit(self, mock_cli_class):
        """Test main function with delete and exit action."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'delete-and-exit', '--job_id', 'test_job_123']):
            main()

        mock_cli.delete_server_job.assert_called_once_with("test_job_123")

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_delete_and_exit_missing_job_id(self, mock_cli_class):
        """Test main function with delete and exit action missing job_id."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'delete-and-exit']):
            with pytest.raises(SystemExit):
                main()

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_metrics_and_exit(self, mock_cli_class):
        """Test main function with metrics and exit action."""
        mock_cli = Mock()
        mock_cli.evaluate_results_locally.return_value = Path("results")
        mock_cli.find_invalid_ontologies.return_value = []
        mock_cli_class.return_value = mock_cli
        
        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'metrics-and-exit', '--local_metrics_ontologies_ids', '1', '2']):
            main()
        
        mock_cli.evaluate_results_locally.assert_called_once_with([1, 2])

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_metrics_and_exit_missing_ontologies(self, mock_cli_class):
        """Test main function with metrics and exit action missing ontologies."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'metrics-and-exit']):
            with pytest.raises(SystemExit):
                main()

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_download_metrics_and_exit(self, mock_cli_class):
        """Test main function with download-metrics-and-exit action."""
        mock_cli = Mock()
        mock_cli.poll_job_to_completion.return_value = None
        mock_cli.extract_server_job_ontologies.return_value = [1, 2]
        mock_cli.download_job_results.return_value = None
        mock_cli.evaluate_results_locally.return_value = Path("results")
        mock_cli.find_invalid_ontologies.return_value = []
        mock_cli_class.return_value = mock_cli
        
        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'download-metrics-and-exit', '--job_id', 'test_job_123']):
            main()
        
        mock_cli.poll_job_to_completion.assert_called_once_with("test_job_123")
        # Note: extract_server_job_ontologies is called twice in the actual code
        assert mock_cli.extract_server_job_ontologies.call_count == 2
        mock_cli.download_job_results.assert_called_once_with("test_job_123", [1, 2])
        mock_cli.evaluate_results_locally.assert_called_once_with([1, 2])

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_download_metrics_and_exit_missing_job_id(self, mock_cli_class):
        """Test main function with download-metrics-and-exit action missing job_id."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'download-metrics-and-exit']):
            with pytest.raises(SystemExit):
                main()

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_list_benchmark_jobs_and_exit(self, mock_cli_class):
        """Test main function with list-benchmark-jobs-and-exit action."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'list-benchmark-jobs-and-exit']):
            main()

        mock_cli.print_server_jobs.assert_called_once_with(False)

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_create_download_metrics_delete(self, mock_cli_class):
        """Test main function with create-download-metrics-delete action."""
        mock_cli = Mock()
        mock_cli.find_invalid_ontologies.return_value = []
        mock_cli.create_benchmark.return_value = "test_job_123"
        mock_cli.poll_job_to_completion.return_value = None
        mock_cli.extract_server_job_ontologies.return_value = [1, 2]
        mock_cli.download_job_results.return_value = None
        mock_cli.evaluate_results_locally.return_value = Path("results")
        mock_cli.delete_server_job.return_value = None
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'create-download-metrics-delete', '--ontologies_ids', '1', '2']):
            main()

        mock_cli.create_benchmark.assert_called_once_with([1, 2])
        mock_cli.poll_job_to_completion.assert_called_once_with("test_job_123")
        mock_cli.download_job_results.assert_called_once_with("test_job_123", [1, 2])
        mock_cli.evaluate_results_locally.assert_called_once_with([1, 2])
        mock_cli.delete_server_job.assert_called_once_with("test_job_123")

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_create_download_metrics_delete_missing_ontologies(self, mock_cli_class):
        """Test main function with create-download-metrics-delete action missing ontologies."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'create-download-metrics-delete']):
            with pytest.raises(SystemExit):
                main()

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_create_download_metrics(self, mock_cli_class):
        """Test main function with create-download-metrics action."""
        mock_cli = Mock()
        mock_cli.find_invalid_ontologies.return_value = []
        mock_cli.create_benchmark.return_value = "test_job_123"
        mock_cli.poll_job_to_completion.return_value = None
        mock_cli.extract_server_job_ontologies.return_value = [1, 2]
        mock_cli.download_job_results.return_value = None
        mock_cli.evaluate_results_locally.return_value = Path("results")
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'create-download-metrics', '--ontologies_ids', '1', '2']):
            main()

        mock_cli.create_benchmark.assert_called_once_with([1, 2])
        mock_cli.poll_job_to_completion.assert_called_once_with("test_job_123")
        mock_cli.download_job_results.assert_called_once_with("test_job_123", [1, 2])
        mock_cli.evaluate_results_locally.assert_called_once_with([1, 2])

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_create_download_metrics_missing_ontologies(self, mock_cli_class):
        """Test main function with create-download-metrics action missing ontologies."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'create-download-metrics']):
            with pytest.raises(SystemExit):
                main()

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_with_detailed_flag(self, mock_cli_class):
        """Test main function with detailed flag."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'list-benchmark-jobs-and-exit', '--detailed']):
            main()

        mock_cli.print_server_jobs.assert_called_once_with(True)

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_with_custom_server_url(self, mock_cli_class):
        """Test main function with custom server URL."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli
        
        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'list-ontologies-ids-and-exit', '--server_url', 'http://custom:8000/api/']):
            main()
        
        mock_cli_class.assert_called_once()
        call_args = mock_cli_class.call_args
        # Check that the server URL was passed as the second positional argument
        assert call_args[0][1] == 'http://custom:8000/api/'

    @patch('cleverswarm_python_client.cswarm_benchmark_client.BenchmarkCLI')
    def test_main_with_username_and_token(self, mock_cli_class):
        """Test main function with username and token."""
        mock_cli = Mock()
        mock_cli_class.return_value = mock_cli

        with patch('sys.argv', ['cswarm_benchmark_client.py', '--action', 'list-ontologies-ids-and-exit', '--username', 'testuser', '--token', 'testtoken']):
            main()

        mock_cli_class.assert_called_once()
        call_args = mock_cli_class.call_args
        assert call_args.kwargs['username'] == 'testuser'
        assert call_args.kwargs['token'] == 'testtoken'