import pytest
from unittest.mock import Mock, patch
from click.exceptions import Exit as ClickExit
from cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs import BigqueryOrgs
import typer


class TestBigqueryOrgs:
    @pytest.fixture
    def bigquery_orgs(self):
        return BigqueryOrgs()

    @patch("cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.bigquery.Client")
    def test_get_bigquery_client_env_creds(self, mock_client):
        """Test get_bigquery_client with environment credentials"""
        with patch.dict(
            "os.environ",
            {"GOOGLE_APPLICATION_CREDENTIALS_JSON": '{"type": "service_account"}'},
        ):
            with patch(
                "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs."
                "service_account.Credentials.from_service_account_info"
            ) as mock_creds:
                mock_creds.return_value = Mock()
                BigqueryOrgs.get_bigquery_client()
                mock_client.assert_called_once()

    @patch("cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.bigquery.Client")
    def test_get_bigquery_client_default_creds(self, mock_client):
        """Test get_bigquery_client with default credentials"""
        with patch.dict("os.environ", {}, clear=True):
            BigqueryOrgs.get_bigquery_client()
            mock_client.assert_called_once_with(credentials=None)

    def test_print_table(self):
        """Test print_table method"""
        headers = ["Name", "Age"]
        rows = [["John", "25"], ["Jane", "30"]]

        with patch("builtins.print") as mock_print:
            BigqueryOrgs.print_table(headers, rows)
            assert mock_print.call_count == 4

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs."
        "get_bigquery_client"
    )
    def test_list_all_success(self, mock_get_client, bigquery_orgs):
        """Test list_all method with successful query"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Create a proper mock for the query result
        mock_row1 = Mock()
        mock_row1.org_id = "CXO001"
        mock_row1.org_code = "CODE1"
        mock_row1.org_type = "TYPE1"
        mock_row1.company_name = "Company1"

        mock_row2 = Mock()
        mock_row2.org_id = "CXO002"
        mock_row2.org_code = "CODE2"
        mock_row2.org_type = "TYPE2"
        mock_row2.company_name = "Company2"

        mock_results = Mock()
        mock_results.total_rows = 2
        mock_results.__iter__ = Mock(return_value=iter([mock_row1, mock_row2]))

        mock_query_job = Mock()
        mock_query_job.result.return_value = mock_results
        mock_client.query.return_value = mock_query_job

        with patch.object(BigqueryOrgs, "print_table") as mock_print:
            bigquery_orgs.list_all()
            mock_print.assert_called_once()

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs."
        "get_bigquery_client"
    )
    def test_list_all_empty(self, mock_get_client, bigquery_orgs):
        """Test list_all method with empty results"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_results = Mock()
        mock_results.total_rows = 0

        mock_query_job = Mock()
        mock_query_job.result.return_value = mock_results
        mock_client.query.return_value = mock_query_job

        with patch("builtins.print") as mock_print:
            bigquery_orgs.list_all()
            mock_print.assert_called_with("No organizations found")

    def test_placeholder_methods(self, bigquery_orgs):
        """Test that placeholder methods print expected messages"""
        with patch("builtins.print") as mock_print:
            bigquery_orgs.accounts_list("CXO00118")
            mock_print.assert_called_with(
                "Listing accounts for org CXO00118 in BigQuery (not implemented)"
            )

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_list_success(self, mock_get_client, bigquery_orgs):
        """Test provision_list method with successful queries"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock subscriptions query results
        mock_sub_row = Mock()
        mock_sub_row.cdu_id = "CDU001"
        mock_sub_row.service_type = "subscription"
        mock_sub_row.validFrom = Mock()
        mock_sub_row.validFrom.strftime.return_value = "2024-01-01"
        mock_sub_row.validTo = Mock()
        mock_sub_row.validTo.strftime.return_value = "2024-12-31"

        mock_subs_job = Mock()
        mock_subs_job.result.return_value = [mock_sub_row]

        # Mock views query results
        mock_view_row = Mock()
        mock_view_row.view_name = "CDU001__view_name"
        mock_view_row.creation_time = Mock()
        mock_view_row.creation_time.strftime.return_value = "2024-01-01"

        mock_views_job = Mock()
        mock_views_job.result.return_value = [mock_view_row]

        # Set up client to return different jobs for different queries
        def query_side_effect(query):
            if "subscriptions" in query:
                return mock_subs_job
            elif "INFORMATION_SCHEMA" in query:
                return mock_views_job
            return Mock()

        mock_client.query.side_effect = query_side_effect

        with patch.object(BigqueryOrgs, "print_table") as mock_print_table:
            bigquery_orgs.provision_list("CXO00118")

            # Should print tables for both subscriptions and views
            assert mock_print_table.call_count == 2

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_list_no_data(self, mock_get_client, bigquery_orgs):
        """Test provision_list method with no subscriptions or views"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock empty results
        mock_subs_job = Mock()
        mock_subs_job.result.return_value = []  # No subscriptions

        mock_views_job = Mock()
        mock_views_job.result.return_value = []  # No views

        def query_side_effect(query):
            if "subscriptions" in query:
                return mock_subs_job
            elif "INFORMATION_SCHEMA" in query:
                return mock_views_job
            return Mock()

        mock_client.query.side_effect = query_side_effect

        with patch("builtins.print") as mock_print:
            bigquery_orgs.provision_list("CXO00118")

            # Should print summary with zero counts
            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            error_messages = [str(call) for call in calls]
            assert any("Subscribed CDUs: 0" in msg for msg in error_messages)
            assert any("Provisioned Views: 0" in msg for msg in error_messages)

    def test_provision_list_invalid_org_id(self, bigquery_orgs):
        """Test provision_list with invalid organization ID"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.provision_list("INVALID_ID")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            error_messages = [str(call) for call in calls]
            assert any("must start with 'CXO'" in msg for msg in error_messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_list_table_not_found(self, mock_get_client, bigquery_orgs):
        """Test entitlements_list when table doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock BigQuery error for non-existent table
        from google.api_core.exceptions import NotFound

        mock_client.query.side_effect = NotFound("Table not found")

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_list("CXO00118")

            # Should provide helpful error message
            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            error_messages = [str(call) for call in calls]
            # Check for any error message containing relevant keywords
            assert any(
                "error" in msg.lower() or "not found" in msg.lower()
                for msg in error_messages
            )

    def test_share_list_basic_functionality(self, bigquery_orgs):
        """Test that share_list method exists and can be called without crashing"""
        with patch("builtins.print") as mock_print:
            # Mock the get_bigquery_client method to avoid authentication issues
            with patch.object(bigquery_orgs, "get_bigquery_client") as mock_get_client:
                mock_client = Mock()
                mock_get_client.return_value = mock_client

                # Mock empty query results
                mock_datasets_job = Mock()
                mock_datasets_job.result.return_value = []  # No datasets

                mock_views_job = Mock()
                mock_views_job.result.return_value = []  # No views

                # Set up client to return different jobs for different queries
                def query_side_effect(query):
                    if "INFORMATION_SCHEMA.SCHEMATA" in query:
                        return mock_datasets_job
                    elif "INFORMATION_SCHEMA.TABLES" in query:
                        return mock_views_job
                    return Mock()

                mock_client.query.side_effect = query_side_effect

                # This should not raise an exception due to the try-catch block and mocking
                bigquery_orgs.share_list("CXO00118")

                # The method should print something (either results or error message)
                assert mock_print.called

                # Verify that queries were attempted
                assert (
                    mock_client.query.call_count == 2
                )  # One for datasets, one for views

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_by_org(self, mock_get_client, mock_prompt, bigquery_orgs):
        """Test provision_add for a specific organization"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org query results
        mock_org_row = Mock()
        mock_org_row.id = 123
        mock_org_row.external_id = "CXO00118"
        mock_org_row.org_code = "TEST001"
        mock_org_row.org_type = "Dedicated"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock source project query results
        mock_source_row = Mock()
        mock_source_row.source_project = "crux-data-lake"
        mock_source_row.project = "crux-client-cxo00118"
        mock_source_row.dataset = "client_cxo00118"
        mock_source_row.type = "Dedicated"
        mock_source_job = Mock()
        mock_source_job.result.return_value = [mock_source_row]

        # Mock dataset exists
        mock_client.get_dataset.return_value = Mock()

        # Mock subscriptions view exists
        mock_client.get_table.return_value = Mock()

        # Mock entitled CDUs query
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx12345"
        mock_cdu_row.service_type = "subscription"
        mock_cdu_row.validFrom = None
        mock_cdu_row.validTo = None
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog query (no results - should exit gracefully)
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = []

        # Set up query side effects
        mock_client.query.side_effect = [
            mock_org_job,  # org lookup
            mock_source_job,  # source project lookup
            mock_entitled_job,  # entitled CDUs
            mock_catalog_job,  # catalog lookup (empty)
        ]

        with patch("builtins.print") as mock_print:
            # FIXED: Call with org_id and dry_run (view_name defaults to None)
            bigquery_orgs.provision_add(org_id="CXO00118", dry_run=True)

            # Should print that no tables found in catalog
            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("No tables found in catalog" in msg for msg in messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_delete_success(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test provision_delete with successful deletion"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_prompt.return_value = "DELETE"

        bigquery_orgs.provision_delete("CXO00118", "test_view")

        # Verify the view deletion was attempted
        mock_client.delete_table.assert_called_once()
        expected_view_id = "crux-client-cxo00118.client_cxo00118.test_view"
        mock_client.delete_table.assert_called_with(expected_view_id)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_list_simplified_success(self, mock_get_client, bigquery_orgs):
        """Test simplified list method with successful query"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_row = Mock()
        mock_row.org_id = "CXO00118"
        mock_row.org_code = "TEST001"
        mock_row.org_type = "ENTERPRISE"
        mock_row.company_name = "Test Company"

        mock_results = Mock()
        mock_results.total_rows = 1
        mock_results.__iter__ = Mock(return_value=iter([mock_row]))

        mock_query_job = Mock()
        mock_query_job.result.return_value = mock_results
        mock_client.query.return_value = mock_query_job

        with patch.object(BigqueryOrgs, "print_table") as mock_print:
            bigquery_orgs.list("CXO00118")
            mock_print.assert_called_once()

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_list_simplified_not_found(self, mock_get_client, bigquery_orgs):
        """Test simplified list method when organization not found"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_results = Mock()
        mock_results.total_rows = 0

        mock_query_job = Mock()
        mock_query_job.result.return_value = mock_results
        mock_client.query.return_value = mock_query_job

        with patch("builtins.print") as mock_print:
            bigquery_orgs.list("CXO00118")
            mock_print.assert_called_with("Organization CXO00118 not found")

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_delete_simplified_not_found(self, mock_get_client, bigquery_orgs):
        """Test simplified delete method when organization not found"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_check_job = Mock()
        mock_check_job.result.return_value = []
        mock_client.query.return_value = mock_check_job

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.delete("CXO00118")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            error_messages = [str(call) for call in calls]
            assert any("not found" in msg for msg in error_messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_add_simplified_success(self, mock_get_client, mock_prompt, bigquery_orgs):
        """Test simplified add method with successful insertion"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_check_job = Mock()
        mock_check_job.result.return_value = []

        mock_insert_job1 = Mock()
        mock_insert_job1.result.return_value = None

        mock_insert_job2 = Mock()
        mock_insert_job2.result.return_value = None

        # Need 3 queries: check, insert pa_identity, insert org_projects
        mock_client.query.side_effect = [
            mock_check_job,
            mock_insert_job1,
            mock_insert_job2,
        ]

        # Updated prompts to match new validation and auto-generation
        mock_prompt.side_effect = [
            "CXO00118",  # org_code (with default)
            "Dedicated",  # org_type (valid - exits loop)
            "Test Company",  # company_name
            "crux-data-lake",  # source_project (valid - exits loop, uses default)
            "ADD",  # confirmation
        ]

        with patch.object(BigqueryOrgs, "print_table"):
            with patch("builtins.print") as mock_print:
                bigquery_orgs.add("CXO00118")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                error_messages = [str(call) for call in calls]
                assert any("Successfully added" in msg for msg in error_messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_add_simplified_cancelled(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test simplified add method when user cancels"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_check_job = Mock()
        mock_check_job.result.return_value = []
        mock_client.query.return_value = mock_check_job

        # Updated prompts to match new validation and auto-generation
        mock_prompt.side_effect = [
            "CXO00118",  # org_code
            "Dedicated",  # org_type (valid - exits loop)
            "Test Company",  # company_name
            "crux-data-lake",  # source_project (valid - exits loop)
            "CANCEL",  # confirmation (not "ADD")
        ]

        with patch.object(BigqueryOrgs, "print_table"):
            with patch("builtins.print") as mock_print:
                with pytest.raises(ClickExit):
                    bigquery_orgs.add("CXO00118")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                error_messages = [str(call) for call in calls]
                assert any("Add operation cancelled" in msg for msg in error_messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_delete_simplified_success(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test simplified delete method with successful deletion"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock underlying tables check
        mock_underlying_row = Mock()
        mock_underlying_row.table_name = "org_projects"
        mock_underlying_row.external_id = "CXO00118"

        mock_underlying_check = Mock()
        mock_underlying_check.result.return_value = [mock_underlying_row]

        # Mock view check
        mock_view_row = Mock()
        mock_view_row.external_id = "CXO00118"
        mock_view_row.org_code = "TEST001"
        mock_view_row.org_type = "ENTERPRISE"
        mock_view_row.company_name = "Test Company"

        mock_view_check = Mock()
        mock_view_check.result.return_value = [mock_view_row]

        # Mock delete jobs
        mock_delete_job1 = Mock()
        mock_delete_job1.result.return_value = None

        mock_delete_job2 = Mock()
        mock_delete_job2.result.return_value = None

        # Need 4 queries: underlying check, view check, delete org_projects, delete pa_identity
        mock_client.query.side_effect = [
            mock_underlying_check,
            mock_view_check,
            mock_delete_job1,
            mock_delete_job2,
        ]

        mock_prompt.side_effect = ["DELETE", "CONFIRM"]

        with patch.object(BigqueryOrgs, "print_table"):
            with patch("builtins.print") as mock_print:
                bigquery_orgs.delete("CXO00118")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                error_messages = [str(call) for call in calls]
                assert any("Successfully deleted" in msg for msg in error_messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_list_routes_to_org(self, mock_get_client, bigquery_orgs):
        """Test provision_list routes to _provision_list_by_org when org_id provided"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock queries for _provision_list_by_org
        mock_subs_job = Mock()
        mock_subs_job.result.return_value = []

        mock_views_job = Mock()
        mock_views_job.result.return_value = []

        mock_client.query.side_effect = [mock_subs_job, mock_views_job]

        with patch.object(bigquery_orgs, "_provision_list_by_org") as mock_by_org:
            bigquery_orgs.provision_list(org_id="CXO00118")
            mock_by_org.assert_called_once_with("CXO00118")

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_list_routes_to_cdu(self, mock_get_client, bigquery_orgs):
        """Test provision_list routes to _provision_list_by_cdu when cdu_id provided"""
        with patch.object(bigquery_orgs, "_provision_list_by_cdu") as mock_by_cdu:
            bigquery_orgs.provision_list(cdu_id="cx12345")
            mock_by_cdu.assert_called_once_with("cx12345")

    def test_provision_list_requires_parameter(self, bigquery_orgs):
        """Test provision_list fails when neither org_id nor cdu_id provided"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.provision_list()

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any(
                "Either --org_id or --cdu_id must be provided" in msg
                for msg in messages
            )

    def test_provision_list_rejects_both_parameters(self, bigquery_orgs):
        """Test provision_list fails when both org_id and cdu_id provided"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.provision_list(org_id="CXO00118", cdu_id="cx12345")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any(
                "Please provide either --org_id or --cdu_id, not both" in msg
                for msg in messages
            )

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_routes_to_org(self, mock_get_client, bigquery_orgs):
        """Test provision_add routes to _provision_add_by_org when org_id provided"""
        with patch.object(bigquery_orgs, "_provision_add_by_org") as mock_by_org:
            # FIXED: Just pass org_id (view_name defaults to None)
            bigquery_orgs.provision_add(org_id="CXO00118")
            # FIXED: Check it's called with (org_id, None, False)
            mock_by_org.assert_called_once_with("CXO00118", None, False)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_routes_to_cdu(self, mock_get_client, bigquery_orgs):
        """Test provision_add routes to _provision_add_by_cdu when cdu_id provided"""
        with patch.object(bigquery_orgs, "_provision_add_by_cdu") as mock_by_cdu:
            bigquery_orgs.provision_add(cdu_id="cx12345")
            mock_by_cdu.assert_called_once_with("cx12345", False)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_by_cdu_no_orgs(self, mock_get_client, bigquery_orgs):
        """Test _provision_add_by_cdu when no orgs are entitled"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx12345"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock entitlements lookup (empty)
        mock_ent_job = Mock()
        mock_ent_job.result.return_value = []

        mock_client.query.side_effect = [mock_cdu_job, mock_ent_job]

        with patch("builtins.print") as mock_print:
            bigquery_orgs._provision_add_by_cdu("cx12345", dry_run=False)

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("No organizations are entitled" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_list_by_cdu_not_found(self, mock_get_client, bigquery_orgs):
        """Test _provision_list_by_cdu when CDU doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = []
        mock_client.query.return_value = mock_cdu_job

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs._provision_list_by_cdu("invalid_cdu")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("CDU invalid_cdu not found" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_by_org_full_flow(self, mock_get_client, bigquery_orgs):
        """Test full _provision_add_by_org workflow with view creation"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org query
        mock_org_row = Mock()
        mock_org_row.id = 123
        mock_org_row.external_id = "CXO00118"
        mock_org_row.org_code = "TEST001"
        mock_org_row.org_type = "Dedicated"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock source project query
        mock_source_row = Mock()
        mock_source_row.source_project = "crux-data-lake"
        mock_source_row.project = "crux-client-cxo00118"
        mock_source_row.dataset = "client_cxo00118"
        mock_source_row.type = "Dedicated"
        mock_source_job = Mock()
        mock_source_job.result.return_value = [mock_source_row]

        # Mock dataset and table existence checks
        mock_client.get_dataset.return_value = Mock()
        mock_client.get_table.side_effect = [
            Mock(),  # subscriptions view exists
            Exception("View not found"),  # First catalog view doesn't exist
        ]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx12345"
        mock_cdu_row.service_type = "subscription"
        mock_cdu_row.validFrom = None
        mock_cdu_row.validTo = None
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog query
        mock_catalog_row = Mock()
        mock_catalog_row.cdu_id = "cx12345"
        mock_catalog_row.table_name = "test_table"
        mock_catalog_row.view_name = "cx12345__test_view"
        mock_catalog_row.dataset_legacy_id = "test_dataset"
        mock_catalog_row.table_dataset = "dataframe_cx12345"
        mock_catalog_row.product_family = "Standard"
        mock_catalog_row.full_table_id = "crux-data-lake.dataframe_cx12345.test_table"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row]

        # Mock crux schema query
        mock_schema_row = Mock()
        mock_schema_row.cdu_id = "cx12345"
        mock_schema_row.table_name = "test_table"
        mock_schema_row.field_name = "test_field"
        mock_schema_row.data_type = "STRING"
        mock_schema_row.ord = 1
        mock_schema_row.description = "Test description"
        mock_schema_row.cxref = None
        mock_schema_row.pkey = False
        mock_schema_job = Mock()
        mock_schema_job.result.return_value = [mock_schema_row]

        # Mock table schema query
        mock_table_schema_row = Mock()
        mock_table_schema_row.table_name = "test_table"
        mock_table_schema_row.ordinal_position = 1
        mock_table_schema_row.column_name = "test_field"
        mock_table_schema_row.data_type = "STRING"
        mock_table_schema_job = Mock()
        mock_table_schema_job.result.return_value = [mock_table_schema_row]

        # Mock view creation job
        mock_create_view_job = Mock()
        mock_create_view_job.result.return_value = None

        # Mock utility view creation jobs
        mock_util_job = Mock()
        mock_util_job.result.return_value = None

        # Set up query side effects in order
        mock_client.query.side_effect = [
            mock_org_job,  # org lookup
            mock_source_job,  # source project
            mock_entitled_job,  # entitled CDUs
            mock_catalog_job,  # catalog
            mock_schema_job,  # crux schemas
            mock_table_schema_job,  # table schemas
            mock_create_view_job,  # create view
            mock_util_job,  # catalog utility view
            mock_util_job,  # schemas utility view
            mock_util_job,  # validations utility view
        ]

        # Mock get_table for view schema update
        mock_view = Mock()
        mock_view.schema = [Mock(name="test_field", field_type="STRING")]
        mock_client.get_table.side_effect = [
            Mock(),  # subscriptions exists
            Exception("View not found"),  # view doesn't exist yet
            mock_view,  # get view for schema update
        ]

        with patch("builtins.print"):
            bigquery_orgs._provision_add_by_org("CXO00118", None, dry_run=False)

        # Verify view creation was attempted
        assert mock_client.query.call_count >= 7

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_add_by_org_dataset_creation(
        self, mock_get_client, bigquery_orgs
    ):
        """Test _provision_add_by_org creates dataset when it doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Setup basic mocks
        mock_org_row = Mock()
        mock_org_row.external_id = "CXO00118"
        mock_org_row.org_code = "TEST001"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        mock_source_row = Mock()
        mock_source_row.source_project = "crux-data-lake"
        mock_source_row.project = "crux-client-cxo00118"
        mock_source_row.dataset = "client_cxo00118"
        mock_source_job = Mock()
        mock_source_job.result.return_value = [mock_source_row]

        # Dataset doesn't exist
        mock_client.get_dataset.side_effect = Exception("Dataset not found")
        mock_client.create_dataset.return_value = Mock()

        # Mock entitled CDUs (empty - will exit early)
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = []

        mock_client.query.side_effect = [
            mock_org_job,
            mock_source_job,
            mock_entitled_job,
        ]

        with patch("builtins.print"):
            with pytest.raises(ClickExit):
                bigquery_orgs._provision_add_by_org("CXO00118", None, dry_run=False)

        # Verify dataset creation was attempted
        mock_client.create_dataset.assert_called_once()


# ============================================================================
# Tests for provision_missing_views method
# ============================================================================


class TestProvisionMissingViews:
    @pytest.fixture
    def bigquery_orgs(self):
        return BigqueryOrgs()

    def test_provision_missing_views_no_org_id(self, bigquery_orgs):
        """Test provision_missing_views fails when org_id is not provided"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id=None)

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("Organization ID is required" in str(msg) for msg in calls)

    def test_provision_missing_views_invalid_format_no_cxo(self, bigquery_orgs):
        """Test provision_missing_views fails when org_id doesn't start with CXO"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="ORG12345")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("must start with 'CXO'" in str(msg) for msg in calls)

    def test_provision_missing_views_invalid_format_non_digits(self, bigquery_orgs):
        """Test provision_missing_views fails when digits aren't present"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="CXOabcde")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("format CXO followed by 5 digits" in str(msg) for msg in calls)

    def test_provision_missing_views_invalid_format_wrong_length(self, bigquery_orgs):
        """Test provision_missing_views fails with wrong length"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="CXO123")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("format CXO followed by 5 digits" in str(msg) for msg in calls)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_org_not_found(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when org not found in org_projects"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        mock_proj_job = Mock()
        mock_proj_job.result.return_value = []
        mock_client.query.return_value = mock_proj_job

        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("not found in org_projects" in str(msg) for msg in calls)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_no_entitlements(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when org has no entitlements"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects query
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"

        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock service_restrictions query (empty)
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = []

        mock_client.query.side_effect = [mock_org_job, mock_entitled_job]

        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            assert any("no entitlements" in str(msg) for msg in calls)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_all_views_present(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when all views are provisioned"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx00001"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog
        mock_catalog_row = Mock()
        mock_catalog_row.view_name = "cx00001__test_view"
        mock_catalog_row.cdu_id = "cx00001"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row]

        # Mock provisioned views (includes all expected + system views)
        mock_view_rows = [
            Mock(table_name="cx00001__test_view"),
            Mock(table_name="subscriptions"),
            Mock(table_name="catalog"),
            Mock(table_name="schemas"),
            Mock(table_name="validations"),
        ]
        mock_views_job = Mock()
        mock_views_job.result.return_value = mock_view_rows

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]
                assert any(
                    "All expected views are provisioned" in msg for msg in messages
                )

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_with_missing_views(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when views are missing"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx00001"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog (2 views)
        mock_catalog_row1 = Mock()
        mock_catalog_row1.view_name = "cx00001__view1"
        mock_catalog_row1.cdu_id = "cx00001"
        mock_catalog_row2 = Mock()
        mock_catalog_row2.view_name = "cx00001__view2"
        mock_catalog_row2.cdu_id = "cx00001"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row1, mock_catalog_row2]

        # Mock provisioned views (only 1 view + system views)
        mock_view_rows = [
            Mock(table_name="cx00001__view1"),
            Mock(table_name="subscriptions"),
            Mock(table_name="catalog"),
            Mock(table_name="schemas"),
            Mock(table_name="validations"),
        ]
        mock_views_job = Mock()
        mock_views_job.result.return_value = mock_view_rows

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table") as mock_table:
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                # Verify print_table was called for missing views
                mock_table.assert_called()
                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]
                assert any("missing views" in msg.lower() for msg in messages)
                assert any("Total missing views: 1" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_with_orphaned_views(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when orphaned views exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx00001"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog
        mock_catalog_row = Mock()
        mock_catalog_row.view_name = "cx00001__view1"
        mock_catalog_row.cdu_id = "cx00001"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row]

        # Mock provisioned views (includes orphaned view)
        mock_view_rows = [
            Mock(table_name="cx00001__view1"),
            Mock(table_name="cx00001__orphaned_view"),
            Mock(table_name="subscriptions"),
            Mock(table_name="catalog"),
            Mock(table_name="schemas"),
            Mock(table_name="validations"),
        ]
        mock_views_job = Mock()
        mock_views_job.result.return_value = mock_view_rows

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]
                assert any("ORPHANED" in msg for msg in messages)
                assert any("Total orphaned views: 1" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_dataset_not_exist(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views when dataset doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx00001"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog
        mock_catalog_row = Mock()
        mock_catalog_row.view_name = "cx00001__view1"
        mock_catalog_row.cdu_id = "cx00001"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row]

        # Mock views query raising "Not found" error
        mock_views_job = Mock()
        mock_views_job.result.side_effect = Exception(
            "Not found: Dataset client_cxo12345"
        )

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]
                assert any("does not exist yet" in msg for msg in messages)
                # When dataset doesn't exist, all expected views (catalog + system) are missing
                assert any("Total missing views: 5" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_multiple_cdus(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views with multiple entitled CDUs"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs (multiple)
        mock_cdu_row1 = Mock()
        mock_cdu_row1.cdu_id = "cx00001"
        mock_cdu_row2 = Mock()
        mock_cdu_row2.cdu_id = "cx00002"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row1, mock_cdu_row2]

        # Mock catalog (views from both CDUs)
        mock_catalog_row1 = Mock()
        mock_catalog_row1.view_name = "cx00001__view1"
        mock_catalog_row1.cdu_id = "cx00001"
        mock_catalog_row2 = Mock()
        mock_catalog_row2.view_name = "cx00002__view1"
        mock_catalog_row2.cdu_id = "cx00002"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row1, mock_catalog_row2]

        # Mock provisioned views (only from cx00001)
        mock_view_rows = [
            Mock(table_name="cx00001__view1"),
            Mock(table_name="subscriptions"),
            Mock(table_name="catalog"),
            Mock(table_name="schemas"),
            Mock(table_name="validations"),
        ]
        mock_views_job = Mock()
        mock_views_job.result.return_value = mock_view_rows

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                # Should show summary
                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]
                assert any("Entitled CDUs: cx00001, cx00002" in msg for msg in messages)
                assert any("Total missing views: 1" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_summary_display(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views displays correct summary"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org_projects
        mock_org_row = Mock()
        mock_org_row.project = "crux-client-cxo12345"
        mock_org_row.dataset = "client_cxo12345"
        mock_org_row.source_project = "crux-data-lake"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock entitled CDUs
        mock_cdu_row = Mock()
        mock_cdu_row.cdu_id = "cx00001"
        mock_entitled_job = Mock()
        mock_entitled_job.result.return_value = [mock_cdu_row]

        # Mock catalog
        mock_catalog_row = Mock()
        mock_catalog_row.view_name = "cx00001__view1"
        mock_catalog_row.cdu_id = "cx00001"
        mock_catalog_job = Mock()
        mock_catalog_job.result.return_value = [mock_catalog_row]

        # Mock provisioned views
        mock_view_rows = [
            Mock(table_name="subscriptions"),
            Mock(table_name="catalog"),
            Mock(table_name="schemas"),
            Mock(table_name="validations"),
        ]
        mock_views_job = Mock()
        mock_views_job.result.return_value = mock_view_rows

        mock_client.query.side_effect = [
            mock_org_job,
            mock_entitled_job,
            mock_catalog_job,
            mock_views_job,
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

                calls = [call_args[0][0] for call_args in mock_print.call_args_list]
                messages = [str(msg) for msg in calls]

                # Verify summary section
                assert any("Expected views:" in msg for msg in messages)
                assert any("Provisioned views:" in msg for msg in messages)
                assert any("Missing views:" in msg for msg in messages)
                assert any("Orphaned views:" in msg for msg in messages)
                assert any("Gap:" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_provision_missing_views_generic_exception(
        self, mock_get_client, bigquery_orgs
    ):
        """Test provision_missing_views handles generic exceptions"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client
        mock_client.query.side_effect = Exception("BigQuery connection error")

        with patch("builtins.print") as mock_print:
            with pytest.raises(typer.Exit):
                bigquery_orgs.provision_missing_views(org_id="CXO12345")

            calls = [call_args[0][0] for call_args in mock_print.call_args_list]
            messages = [str(msg) for msg in calls]
            assert any("Error finding missing views" in msg for msg in messages)

    # Add these tests to your TestBigqueryOrgs class (after the existing tests)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_success(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test entitlements_add with successful addition"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx00001"
        mock_cdu_row.name = "Test CDU"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock duplicate check (no duplicates)
        mock_dup_job = Mock()
        mock_dup_job.result.return_value = []

        # Mock insert job
        mock_insert_job = Mock()
        mock_insert_job.result.return_value = None

        mock_client.query.side_effect = [
            mock_org_job,  # org lookup
            mock_cdu_job,  # cdu lookup
            mock_dup_job,  # duplicate check
            mock_insert_job,  # insert
        ]

        # Mock prompts
        mock_prompt.side_effect = [
            "Full Access",  # service_type
            "2026-01-08",  # service_start_date
            "2027-01-08",  # service_end_date
            "ADD",  # confirmation
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.entitlements_add("CXO12345", "cx00001")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                messages = [str(call) for call in calls]
                assert any("Successfully added entitlement" in msg for msg in messages)

    def test_entitlements_add_invalid_org_id_format(self, bigquery_orgs):
        """Test entitlements_add with invalid org_id format"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_add("INVALID", "cx00001")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("must start with 'CXO'" in msg for msg in messages)

    def test_entitlements_add_invalid_org_id_length(self, bigquery_orgs):
        """Test entitlements_add with wrong length org_id"""
        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_add("CXO123", "cx00001")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("format CXO followed by 5 digits" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_org_not_found(self, mock_get_client, bigquery_orgs):
        """Test entitlements_add when organization doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock empty org lookup
        mock_org_job = Mock()
        mock_org_job.result.return_value = []

        mock_client.query.return_value = mock_org_job

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_add("CXO12345", "cx00001")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("Organization CXO12345 not found" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_cdu_not_found(self, mock_get_client, bigquery_orgs):
        """Test entitlements_add when CDU doesn't exist"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup (success)
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup (not found)
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = []

        mock_client.query.side_effect = [mock_org_job, mock_cdu_job]

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_add("CXO12345", "cx99999")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("CDU cx99999 not found" in msg for msg in messages)

    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_duplicate_entitlement(
        self, mock_get_client, bigquery_orgs
    ):
        """Test entitlements_add when entitlement already exists"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx00001"
        mock_cdu_row.name = "Test CDU"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock duplicate check (found duplicate)
        mock_dup_row = Mock()
        mock_dup_row.cdu_id = "cx00001"
        mock_dup_job = Mock()
        mock_dup_job.result.return_value = [mock_dup_row]

        mock_client.query.side_effect = [mock_org_job, mock_cdu_job, mock_dup_job]

        with patch("builtins.print") as mock_print:
            with pytest.raises(ClickExit):
                bigquery_orgs.entitlements_add("CXO12345", "cx00001")

            calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
            messages = [str(call) for call in calls]
            assert any("already has entitlement" in msg for msg in messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_cancelled(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test entitlements_add when user cancels"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx00001"
        mock_cdu_row.name = "Test CDU"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock duplicate check (no duplicates)
        mock_dup_job = Mock()
        mock_dup_job.result.return_value = []

        mock_client.query.side_effect = [mock_org_job, mock_cdu_job, mock_dup_job]

        # Mock prompts - user cancels
        mock_prompt.side_effect = [
            "Full Access",  # service_type
            "2026-01-08",  # service_start_date
            "2027-01-08",  # service_end_date
            "CANCEL",  # confirmation (not "ADD")
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                with pytest.raises(ClickExit):
                    bigquery_orgs.entitlements_add("CXO12345", "cx00001")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                messages = [str(call) for call in calls]
                assert any("Add operation cancelled" in msg for msg in messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_with_restricted_service(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test entitlements_add with Restricted service type"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx00001"
        mock_cdu_row.name = "Test CDU"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock duplicate check (no duplicates)
        mock_dup_job = Mock()
        mock_dup_job.result.return_value = []

        # Mock insert job
        mock_insert_job = Mock()
        mock_insert_job.result.return_value = None

        mock_client.query.side_effect = [
            mock_org_job,
            mock_cdu_job,
            mock_dup_job,
            mock_insert_job,
        ]

        # Mock prompts with Restricted service type
        mock_prompt.side_effect = [
            "Restricted",  # service_type
            "2026-01-08",  # service_start_date
            "2027-01-08",  # service_end_date
            "ADD",  # confirmation
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                bigquery_orgs.entitlements_add("CXO12345", "cx00001")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                messages = [str(call) for call in calls]
                assert any("Successfully added entitlement" in msg for msg in messages)

    @patch("typer.prompt")
    @patch(
        "cruxctl.command_groups.bigsnow.bigquery.bigquery_orgs.BigqueryOrgs.get_bigquery_client"
    )
    def test_entitlements_add_insert_failure(
        self, mock_get_client, mock_prompt, bigquery_orgs
    ):
        """Test entitlements_add when insert fails"""
        mock_client = Mock()
        mock_get_client.return_value = mock_client

        # Mock org lookup
        mock_org_row = Mock()
        mock_org_row.org_id = "org-internal-id-123"
        mock_org_row.external_id = "CXO12345"
        mock_org_row.company_name = "Test Company"
        mock_org_job = Mock()
        mock_org_job.result.return_value = [mock_org_row]

        # Mock CDU lookup
        mock_cdu_row = Mock()
        mock_cdu_row.code = "cx00001"
        mock_cdu_row.name = "Test CDU"
        mock_cdu_job = Mock()
        mock_cdu_job.result.return_value = [mock_cdu_row]

        # Mock duplicate check (no duplicates)
        mock_dup_job = Mock()
        mock_dup_job.result.return_value = []

        # Mock insert job that fails
        mock_insert_job = Mock()
        mock_insert_job.result.side_effect = Exception("Insert failed")

        mock_client.query.side_effect = [
            mock_org_job,
            mock_cdu_job,
            mock_dup_job,
            mock_insert_job,
        ]

        # Mock prompts
        mock_prompt.side_effect = [
            "Full Access",
            "2026-01-08",
            "2027-01-08",
            "ADD",
        ]

        with patch("builtins.print") as mock_print:
            with patch.object(BigqueryOrgs, "print_table"):
                with pytest.raises(ClickExit):
                    bigquery_orgs.entitlements_add("CXO12345", "cx00001")

                calls = [call_args[0] for call_args, _ in mock_print.call_args_list]
                messages = [str(call) for call in calls]
                assert any("Error inserting entitlement" in msg for msg in messages)
