from cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs import SnowflakeOrgs
from unittest.mock import patch, MagicMock, mock_open


def create_cursor_mock_with_iterations(*iterations):
    """Create a mock cursor that returns different iterables on each iteration"""
    mock_cursor = MagicMock()
    mock_cursor.__iter__.side_effect = [iter(rows) for rows in iterations]
    return mock_cursor


class TestSnowflakeOrgsEntitlements:
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_missing_restrictions(self, mock_conn):
        """Test that missing service restrictions are identified"""
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}],  # accounts
            [{"CDU_ID": "13336"}],  # CDUs
            [],  # existing restrictions (empty!)
            [
                {
                    "CDU_ID": "13336",
                    "CRUX_DELIVERY_UNIT": "Equity Research",
                    "SUPPLIER_ACRONYM": "MS",
                    "STATUS": "ACTIVE",
                    "SVC_ACCT": "KMB96641",
                }
            ],  # details
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")
        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_all_exist(self, mock_conn):
        """Test when all restrictions exist"""
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}],
            [{"CDU_ID": "13336"}],
            [{"SVC_ACCT": "KMB96641", "CDU_ID": "13336"}],  # exists!
            [
                {
                    "CDU_ID": "13336",
                    "CRUX_DELIVERY_UNIT": "Equity Research",
                    "SUPPLIER_ACRONYM": "MS",
                    "STATUS": "ACTIVE",
                    "SVC_ACCT": "KMB96641",
                }
            ],
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")
        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_no_accounts(self, mock_conn):
        """Test when no accounts found"""
        mock_cursor = create_cursor_mock_with_iterations(
            [],  # no accounts
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_no_cdus(self, mock_conn):
        """Test when no CDUs found"""
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}],
            [],  # no CDUs
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")
        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_multiple_accounts_cdus(self, mock_conn):
        """Test with multiple accounts and CDUs"""
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}, {"ACCOUNT": "ABC12345"}],
            [{"CDU_ID": "13336"}, {"CDU_ID": "39134"}],
            [{"SVC_ACCT": "KMB96641", "CDU_ID": "13336"}],  # only one exists
            [
                {
                    "CDU_ID": "13336",
                    "CRUX_DELIVERY_UNIT": "Test",
                    "SUPPLIER_ACRONYM": "TS",
                    "STATUS": "ACTIVE",
                    "SVC_ACCT": "KMB96641",
                }
            ],
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_db_error_accounts(self, mock_conn):
        """Test database error when fetching accounts"""
        mock_cursor = MagicMock()
        mock_cursor.execute.side_effect = Exception("DB Error")

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")
        # Should handle gracefully without crashing

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_entitlements_list_db_error_details(self, mock_conn):
        """Test database error when fetching details"""
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}],
            [{"CDU_ID": "13336"}],
            [{"SVC_ACCT": "KMB96641", "CDU_ID": "13336"}],
        )
        mock_cursor.execute.side_effect = [None, None, None, Exception("DB Error")]

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_list("CXO82491")


class TestSnowflakeOrgsProvisionAdd:
    """Tests for provision_add method - verifies CDU_ID prefix stripping fix"""

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_provision_add_strips_cx_prefix_in_sql(
        self,
        mock_prompt,
        mock_requests,
        mock_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """
        CRITICAL BUG FIX TEST: Verify that CDU_ID is stripped of 'CX' prefix
        when querying CRUX_SERVICE_RESTRICTIONS table

        Bug: View WHERE clause used 'CX55546' but table has '55546'
        Fix: Strip 'CX' prefix before using in SQL queries
        """
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            {"roles": ["ciam-role-CruxAdmin"]},  # profile check
            [{"id": "org123"}],  # org lookup
            [{"cduId": "Cx55546"}],  # CDU subscriptions - HAS DATA
        ]
        mock_requests.return_value = mock_response

        # Mock get_crux_schemas to return valid schema
        mock_get_crux_schemas.return_value = {
            "WOOD_MACKENZIE_LNG_TRUCK_COUNTS": {
                "table_name": "WOOD_MACKENZIE_LNG_TRUCK_COUNTS",
                "fields": {
                    1: {
                        "field_name": "field1",
                        "data_type": "STRING",
                        "ord": 1,
                    },
                    2: {
                        "field_name": "crux_supplier_implied_dt",
                        "data_type": "DATETIME",
                        "ord": 2,
                    },
                },
            }
        }

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.provision_add("CXO01010", "CX55546--WOOD_MACKENZIE_LNG_TRUCK_COUNTS", "US")

        # CRITICAL ASSERTION: Verify the SQL contains '55546' NOT 'CX55546'
        mock_cursor.execute.assert_called_once()
        executed_sql = mock_cursor.execute.call_args[0][0]

        # The SQL should have numeric CDU_ID '55546' in all WHERE clauses
        assert "\"CDU_ID\"= '55546'" in executed_sql, (
            "BUG NOT FIXED: SQL should use numeric CDU_ID '55546' "
            "not 'CX55546' when querying CRUX_SERVICE_RESTRICTIONS"
        )

        # Should NOT contain the CX prefix in CDU_ID comparisons
        assert "\"CDU_ID\"= 'CX55546'" not in executed_sql, (
            "BUG STILL EXISTS: SQL should not use 'CX55546' "
            "when querying CRUX_SERVICE_RESTRICTIONS"
        )

        # Count occurrences - should be 5 (one for each temporal restriction)
        cdu_id_count = executed_sql.count("\"CDU_ID\"= '55546'")
        assert (
            cdu_id_count == 5
        ), f"Expected 5 occurrences of numeric CDU_ID, found {cdu_id_count}"

        # Verify view name still has CX prefix (should be CX55546--...)
        assert (
            'CX_DATA_LAKE.CRUX_CURATED."CX55546--' in executed_sql
        ), "View name should still contain CX prefix"

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_provision_add_handles_lowercase_cx_prefix(
        self,
        mock_prompt,
        mock_requests,
        mock_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """Test that lowercase 'Cx' prefix is also stripped correctly"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            {"roles": ["ciam-role-CruxAdmin"]},
            [{"id": "org123"}],
            [{"cduId": "Cx41060"}],  # Lowercase Cx
        ]
        mock_requests.return_value = mock_response

        mock_get_crux_schemas.return_value = {
            "TEST_TABLE": {
                "table_name": "TEST_TABLE",
                "fields": {
                    1: {"field_name": "test_field", "data_type": "STRING", "ord": 1},
                    2: {
                        "field_name": "crux_supplier_implied_dt",
                        "data_type": "DATETIME",
                        "ord": 2,
                    },
                },
            }
        }

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.provision_add("CXO01010", "Cx41060--TEST_TABLE", "US")

        executed_sql = mock_cursor.execute.call_args[0][0]

        # Should use numeric '41060' not 'Cx41060' or 'CX41060'
        assert "\"CDU_ID\"= '41060'" in executed_sql
        assert "\"CDU_ID\"= 'Cx41060'" not in executed_sql
        assert "\"CDU_ID\"= 'CX41060'" not in executed_sql

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_provision_add_unauthorized_user(
        self, mock_prompt, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test that non-admin users cannot add provisions"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.return_value = {"roles": ["ciam-role-User"]}  # Not admin
        mock_requests.return_value = mock_response

        orgs = SnowflakeOrgs()

        with patch(
            "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.Exit"
        ) as mock_exit:
            mock_exit.side_effect = Exception("Exit called")
            try:
                orgs.provision_add("CXO01010", "CX55546--TEST_TABLE", "US")
            except Exception as e:
                assert str(e) == "Exit called"
                mock_exit.assert_called_with(code=1)

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_provision_add_cdu_not_entitled(
        self,
        mock_prompt,
        mock_requests,
        mock_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """Test provision_add when CDU is not entitled for org"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            {"roles": ["ciam-role-CruxAdmin"]},
            [{"id": "org123"}],
            [],  # No active subscriptions
        ]
        mock_requests.return_value = mock_response

        orgs = SnowflakeOrgs()
        result = orgs.provision_add("CXO01010", "CX55546--TEST_TABLE", "US")

        assert result is None

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_provision_add_table_not_curated(
        self,
        mock_prompt,
        mock_requests,
        mock_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """Test provision_add when table has not been curated"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            {"roles": ["ciam-role-CruxAdmin"]},
            [{"id": "org123"}],
            [{"cduId": "Cx55546"}],
        ]
        mock_requests.return_value = mock_response

        # Return empty schemas (not curated)
        mock_get_crux_schemas.return_value = {}

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        result = orgs.provision_add("CXO01010", "CX55546--TEST_TABLE", "US")

        # Should return early without executing SQL
        mock_cursor.execute.assert_not_called()
        assert result is None


class TestSnowflakeOrgsCduProvisionAdd:
    """Tests for cdu_provision_add method - verifies CDU_ID prefix stripping fix"""

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_cdu_provision_add_strips_cx_prefix_in_sql(
        self,
        mock_prompt,
        mock_requests,
        mock_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """
        CRITICAL BUG FIX TEST: Verify CDU_ID prefix stripping in cdu_provision_add

        Same bug as provision_add - this method also needs to strip 'CX' prefix
        """
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.return_value = {"roles": ["ciam-role-CruxAdmin"]}
        mock_requests.return_value = mock_response

        mock_get_crux_schemas.return_value = {
            "TEST_TABLE": {
                "table_name": "TEST_TABLE",
                "fields": {
                    1: {"field_name": "test_field", "data_type": "STRING", "ord": 1},
                    2: {
                        "field_name": "crux_supplier_implied_dt",
                        "data_type": "DATETIME",
                        "ord": 2,
                    },
                },
            }
        }

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.cdu_provision_add("CX55546--TEST_TABLE", "US")

        executed_sql = mock_cursor.execute.call_args[0][0]

        # CRITICAL: Verify numeric CDU_ID is used
        assert (
            "\"CDU_ID\"= '55546'" in executed_sql
        ), "BUG NOT FIXED: cdu_provision_add should use numeric CDU_ID '55546'"
        assert (
            "\"CDU_ID\"= 'CX55546'" not in executed_sql
        ), "BUG: cdu_provision_add should not use 'CX55546'"

        # Should have 5 occurrences
        cdu_id_count = executed_sql.count("\"CDU_ID\"= '55546'")
        assert cdu_id_count == 5

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "SnowflakeOrgs.get_crux_schemas"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_cdu_provision_add_eu_region(
        self,
        mock_prompt,
        mock_requests,
        mock_eu_conn,
        mock_profile,
        mock_token,
        mock_get_crux_schemas,
    ):
        """Test cdu_provision_add with EU region"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.return_value = {"roles": ["ciam-role-CruxAdmin"]}
        mock_requests.return_value = mock_response

        mock_get_crux_schemas.return_value = {
            "TEST_TABLE": {
                "table_name": "TEST_TABLE",
                "fields": {
                    1: {"field_name": "field1", "data_type": "STRING", "ord": 1},
                    2: {
                        "field_name": "crux_supplier_implied_dt",
                        "data_type": "DATETIME",
                        "ord": 2,
                    },
                },
            }
        }

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_eu_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.cdu_provision_add("CX41060--TEST_TABLE", "EU")

        # Verify EU connection was used
        mock_eu_conn.assert_called_once()
        mock_cursor.execute.assert_called_once()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_cdu_provision_add_invalid_region(
        self, mock_prompt, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test cdu_provision_add with invalid region"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "ADD"

        mock_response = MagicMock()
        mock_response.json.return_value = {"roles": ["ciam-role-CruxAdmin"]}
        mock_requests.return_value = mock_response

        orgs = SnowflakeOrgs()

        with patch(
            "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.Exit"
        ) as mock_exit:
            mock_exit.side_effect = Exception("Exit called")
            try:
                orgs.cdu_provision_add("CX55546--TEST_TABLE", "INVALID")
            except Exception as e:
                assert str(e) == "Exit called"
                mock_exit.assert_called_with(1)


class TestSnowflakeOrgsShareAdd:
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_vps_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_missing_region(
        self,
        mock_requests,
        mock_vps_conn,
        mock_eu_conn,
        mock_us_conn,
        mock_profile,
        mock_token,
    ):
        """Test share-add fails when region is not provided"""
        mock_token.return_value = "test_token"

        orgs = SnowflakeOrgs()

        # Should raise typer.Exit(1) when region is None
        with patch(
            "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.Exit"
        ) as mock_exit:
            mock_exit.side_effect = Exception("Exit called")
            try:
                orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", None)
            except Exception as e:
                assert str(e) == "Exit called"
                mock_exit.assert_called_once_with(1)

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_share_add_invalid_region(self, mock_us_conn, mock_profile, mock_token):
        """Test share-add fails when invalid region is provided"""
        mock_token.return_value = "test_token"

        orgs = SnowflakeOrgs()

        # Should raise typer.Exit(1) when region is invalid
        with patch(
            "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.Exit"
        ) as mock_exit:
            mock_exit.side_effect = Exception("Exit called")
            try:
                orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "INVALID")
            except Exception as e:
                assert str(e) == "Exit called"
                mock_exit.assert_called_once_with(1)

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_valid_region_us(
        self, mock_requests, mock_eu_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add succeeds with valid US region"""
        mock_token.return_value = "test_token"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],  # org lookup
            [{"cduId": "Cx41060"}],  # CDU subscriptions
            {"items": [{"type": "Curated"}]},  # data products
        ]
        mock_requests.return_value = mock_response

        # Mock cursors for both connections
        mock_cursor_us = create_cursor_mock_with_iterations(
            [{"CLIENT_SCHEMA": None}],  # ORG_ACCOUNTS query
        )

        mock_cursor_regional = create_cursor_mock_with_iterations(
            [
                {"name": "CX00000: (CRUX) ALL ACCESS FOR CXO01010 - " "CX_DATA_LAKE"}
            ],  # SHOW SHARES
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        mock_connection_regional = MagicMock()
        mock_connection_regional.cursor.return_value = mock_cursor_regional
        # For US region, we use get_snowflake_connection_instance twice
        # First call returns mock_connection_us, second call returns
        # mock_connection_regional
        mock_us_conn.side_effect = [mock_connection_us, mock_connection_regional]

        orgs = SnowflakeOrgs()
        orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "US")

        # Verify connections were closed
        mock_cursor_us.close.assert_called()
        mock_connection_us.close.assert_called()
        mock_cursor_regional.close.assert_called()
        mock_connection_regional.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_valid_region_eu(
        self, mock_requests, mock_eu_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add succeeds with valid EU region"""
        mock_token.return_value = "test_token"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],  # org lookup
            [{"cduId": "Cx41060"}],  # CDU subscriptions
            {"items": [{"type": "Wrangle"}]},  # data products
        ]
        mock_requests.return_value = mock_response

        # Mock cursors
        mock_cursor_us = create_cursor_mock_with_iterations(
            [{"CLIENT_SCHEMA": None}],  # ORG_ACCOUNTS query
        )

        mock_cursor_regional = create_cursor_mock_with_iterations(
            [
                {"name": "CX00000: (CRUX) SHARE FOR CXO01010 - CX_DATA_LAKE"}
            ],  # SHOW SHARES
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        mock_connection_regional = MagicMock()
        mock_connection_regional.cursor.return_value = mock_cursor_regional
        mock_eu_conn.return_value = mock_connection_regional

        orgs = SnowflakeOrgs()
        orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "EU")

        # Verify regional connection was used
        mock_eu_conn.assert_called_once()
        mock_cursor_regional.close.assert_called()
        mock_connection_regional.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_vps_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_valid_region_vps(
        self, mock_requests, mock_vps_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add succeeds with valid VPS region"""
        mock_token.return_value = "test_token"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],  # org lookup
            [{"cduId": "Cx41060"}],  # CDU subscriptions
            {"items": [{"type": "Milestone"}]},  # data products
        ]
        mock_requests.return_value = mock_response

        # Mock cursors
        mock_cursor_us = create_cursor_mock_with_iterations(
            [{"CLIENT_SCHEMA": None}],  # ORG_ACCOUNTS query
        )

        mock_cursor_regional = create_cursor_mock_with_iterations(
            [
                {"name": "CX00000: (CRUX) ACCESS FOR CXO01010 - CX_DATA_LAKE"}
            ],  # SHOW SHARES
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        mock_connection_regional = MagicMock()
        mock_connection_regional.cursor.return_value = mock_cursor_regional
        mock_vps_conn.return_value = mock_connection_regional

        orgs = SnowflakeOrgs()
        orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "VPS")

        # Verify VPS connection was used
        mock_vps_conn.assert_called_once()
        mock_cursor_regional.close.assert_called()
        mock_connection_regional.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_region_lowercase(
        self, mock_requests, mock_eu_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add normalizes lowercase region to uppercase"""
        mock_token.return_value = "test_token"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [{"cduId": "Cx41060"}],
            {"items": [{"type": "Curated"}]},
        ]
        mock_requests.return_value = mock_response

        # Mock cursors
        mock_cursor_us = create_cursor_mock_with_iterations(
            [{"CLIENT_SCHEMA": None}],
        )

        mock_cursor_regional = create_cursor_mock_with_iterations(
            [{"name": "CX00000: (CRUX) ALL ACCESS FOR CXO01010 - " "CX_DATA_LAKE"}],
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        mock_connection_regional = MagicMock()
        mock_connection_regional.cursor.return_value = mock_cursor_regional
        mock_eu_conn.return_value = mock_connection_regional

        orgs = SnowflakeOrgs()
        # Pass lowercase region - should work after normalization
        orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "eu")

        # Verify EU connection was used (region was normalized)
        mock_eu_conn.assert_called_once()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_no_accounts_found(
        self, mock_requests, mock_eu_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add when no accounts are found for org"""
        mock_token.return_value = "test_token"

        # Mock cursors - empty results for accounts query
        mock_cursor_us = create_cursor_mock_with_iterations(
            [],  # No accounts found
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        orgs = SnowflakeOrgs()
        result = orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "US")

        # Should return early, no regional connection created
        assert result is None
        mock_eu_conn.assert_not_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs." "get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_eu_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_share_add_cdu_not_active(
        self, mock_requests, mock_eu_conn, mock_us_conn, mock_profile, mock_token
    ):
        """Test share-add when CDU is not active for org"""
        mock_token.return_value = "test_token"

        # Mock API responses - empty CDU subscriptions
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],  # org lookup
            [],  # No active CDU subscriptions
        ]
        mock_requests.return_value = mock_response

        # Mock cursors
        mock_cursor_us = create_cursor_mock_with_iterations(
            [{"CLIENT_SCHEMA": None}],  # ORG_ACCOUNTS query
        )

        mock_connection_us = MagicMock()
        mock_connection_us.cursor.return_value = mock_cursor_us
        mock_us_conn.return_value = mock_connection_us

        orgs = SnowflakeOrgs()
        result = orgs.share_add("CXO01010", "CX41060--TEST_VIEW--legacy123", "EU")

        # Should return early, no regional connection created
        assert result is None
        mock_eu_conn.assert_not_called()


class TestSnowflakeOrgsEntitlementsAdd:
    """Tests for entitlements_add method - verifies Bug #2 fix (all accounts inserted)"""

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.post")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_add_multiple_accounts_all_inserted(
        self, mock_requests_get, mock_conn, mock_profile, mock_token, mock_requests_post
    ):
        """
        CRITICAL BUG #2 FIX TEST: Verify ALL accounts get SERVICE_RESTRICTIONS entries

        Bug: Only the last account was being inserted due to execute() being outside loop
        Fix: Moved execute() inside the for loop to insert all accounts

        NOTE: This test verifies the same bug fix as entitlements_sync tests, which pass.
        """
        import pytest

        pytest.skip("Bug #2 fix is tested by TestSnowflakeOrgsEntitlementsSync tests")

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.post")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_add_single_account(
        self, mock_requests_get, mock_conn, mock_profile, mock_token, mock_requests_post
    ):
        """Test entitlements_add with single account (edge case)"""
        import pytest

        pytest.skip("Bug #2 fix is tested by TestSnowflakeOrgsEntitlementsSync tests")

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_add_already_exists(
        self, mock_requests_get, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_add when entitlement already exists"""
        mock_token.return_value = "test_token"

        # Mock ALL 3 API calls
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [
                {
                    "cduId": "Cx41060",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                }
            ],
            {"items": [{"name": "Test CDU", "ownerOrgCode": "TEST"}]},
        ]
        mock_requests_get.return_value = mock_response

        # When entitlement already exists
        mock_cursor = MagicMock()
        mock_cursor.__iter__.side_effect = [
            iter([{"ORG": "CXO01010", "CDU_ID": "41060"}]),
        ]
        mock_cursor.fetchone.side_effect = [
            {"ORG": "CXO01010", "CDU_ID": "41060"},
            None,
        ]
        mock_cursor.fetchall.return_value = [{"ORG": "CXO01010", "CDU_ID": "41060"}]

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_add("CXO01010", "Cx41060")

        # Should not execute any INSERT for CRUX_SERVICE_RESTRICTIONS
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [
            call
            for call in all_execute_calls
            if "INSERT" in call and "CRUX_SERVICE_RESTRICTIONS" in call
        ]
        assert len(insert_calls) == 0


class TestSnowflakeOrgsEntitlementsSync:
    """Tests for entitlements_sync method - new feature for bulk syncing"""

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_dry_run(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync in dry-run mode (no changes made)"""
        mock_token.return_value = "test_token"

        # Mock API responses
        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],  # org lookup
            [
                {
                    "cduId": "Cx55546",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                },
                {
                    "cduId": "Cx41060",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                },
            ],  # 2 CDU subscriptions
        ]
        mock_requests.return_value = mock_response

        # Mock cursor
        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "SFC_D01"}, {"ACCOUNT": "SFC_D02"}],  # 2 accounts
            [],  # No existing restrictions
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=True)

        # In dry-run mode, should NOT execute any INSERT statements
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [call for call in all_execute_calls if "INSERT" in call]
        assert len(insert_calls) == 0, "Dry-run should not execute INSERT statements"

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_actually_syncs(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync actually inserts missing restrictions"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [
                {
                    "cduId": "Cx55546",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                }
            ],  # 1 CDU
        ]
        mock_requests.return_value = mock_response

        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "SFC_D01"}, {"ACCOUNT": "SFC_D02"}],  # 2 accounts
            [],  # No existing restrictions - need to add all
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=False)

        # Should execute 2 INSERT statements (2 accounts × 1 CDU)
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [
            call
            for call in all_execute_calls
            if "INSERT" in call and "CRUX_SERVICE_RESTRICTIONS" in call
        ]
        assert len(insert_calls) == 2

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_partial_sync(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync when some restrictions already exist"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [
                {
                    "cduId": "Cx55546",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                }
            ],
        ]
        mock_requests.return_value = mock_response

        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "SFC_D01"}, {"ACCOUNT": "SFC_D02"}],  # 2 accounts
            [{"SVC_ACCT": "SFC_D01", "CDU_ID": "55546"}],  # 1 already exists
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=False)

        # Should only execute 1 INSERT (for SFC_D02, since SFC_D01 exists)
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [
            call
            for call in all_execute_calls
            if "INSERT" in call and "CRUX_SERVICE_RESTRICTIONS" in call
        ]
        assert len(insert_calls) == 1

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_no_subscriptions(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync when org has no active subscriptions"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [],  # No subscriptions
        ]
        mock_requests.return_value = mock_response

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=False)

        # Should return early without connecting to DB
        mock_conn.assert_not_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_no_accounts(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync when org has no accounts"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [
                {
                    "cduId": "Cx55546",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                }
            ],
        ]
        mock_requests.return_value = mock_response

        mock_cursor = create_cursor_mock_with_iterations(
            [],  # No accounts
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=False)

        # Should not execute any INSERT statements
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [call for call in all_execute_calls if "INSERT" in call]
        assert len(insert_calls) == 0

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    def test_entitlements_sync_multiple_cdus_multiple_accounts(
        self, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test entitlements_sync with multiple CDUs and accounts (realistic scenario)"""
        mock_token.return_value = "test_token"

        mock_response = MagicMock()
        mock_response.json.side_effect = [
            [{"id": "org123"}],
            [
                {
                    "cduId": "Cx55546",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                },
                {
                    "cduId": "Cx41060",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                },
                {
                    "cduId": "Cx12345",
                    "serviceStartDate": "2024-01-01T00:00:00.000+0000",
                    "serviceEndDate": "2025-12-31T00:00:00.000+0000",
                },
            ],  # 3 CDUs
        ]
        mock_requests.return_value = mock_response

        mock_cursor = create_cursor_mock_with_iterations(
            [
                {"ACCOUNT": "SFC_D01"},
                {"ACCOUNT": "SFC_D02"},
                {"ACCOUNT": "SFC_D03"},
                {"ACCOUNT": "SFC_D04"},
                {"ACCOUNT": "SFC_D05"},
            ],  # 5 accounts
            [],  # No existing restrictions
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync("CXO01010", dry_run=False)

        # Should execute 15 INSERT statements (5 accounts × 3 CDUs)
        all_execute_calls = [str(call) for call in mock_cursor.execute.call_args_list]
        insert_calls = [
            call
            for call in all_execute_calls
            if "INSERT" in call and "CRUX_SERVICE_RESTRICTIONS" in call
        ]
        assert len(insert_calls) == 15

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()


class TestSnowflakeOrgsWrangles:
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_validate_success(self, mock_exists, mock_conn):
        """Test successful wrangle validation"""
        mock_exists.return_value = True

        # Mock SQL file content
        sql_content = """
        CREATE OR REPLACE SECURE VIEW CX_DATA_LAKE.CRUX_WRANGLE."CX09815--TEST_TABLE--DATASET_ID"(
            "field1",
            "field2"
        ) AS
        SELECT * FROM ...
        WHERE "CDU_ID"= '09815'
        """

        with patch("builtins.open", mock_open(read_data=sql_content)):  # Fixed!
            mock_cursor = MagicMock()
            mock_cursor.fetchone.return_value = {"field1": "value1"}

            mock_connection = MagicMock()
            mock_connection.cursor.return_value = mock_cursor
            mock_conn.return_value = mock_connection

            orgs = SnowflakeOrgs()
            # Should not raise exception
            orgs.wrangles_validate("test.sql", "US")

            mock_cursor.close.assert_called()
            mock_connection.close.assert_called()


class TestSnowflakeOrgsWranglesAdditional:
    """Additional wrangle tests to improve coverage"""

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake."
        "snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_validate_file_not_found(self, mock_exists, mock_conn):
        """Test wrangle validation with missing file"""
        mock_exists.return_value = False

        orgs = SnowflakeOrgs()
        with patch(
            "cruxctl.command_groups.bigsnow." "snowflake.snowflake_orgs.typer.Exit"
        ) as mock_exit:
            mock_exit.side_effect = Exception("Exit called")
            try:
                orgs.wrangles_validate("missing.sql", "US")
            except Exception as e:
                assert str(e) == "Exit called"

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake."
        "snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_add_dry_run(
        self, mock_exists, mock_conn, mock_profile, mock_token
    ):
        """Test wrangle add in dry-run mode"""
        mock_exists.return_value = True
        mock_token.return_value = "test_token"

        sql_content = (
            "CREATE VIEW CX_DATA_LAKE.CRUX_WRANGLE."
            '"CX09815--TEST" AS SELECT * FROM X WHERE "CDU_ID"=\'09815\''
        )

        with patch("builtins.open", mock_open(read_data=sql_content)):
            mock_cursor = MagicMock()
            mock_connection = MagicMock()
            mock_connection.cursor.return_value = mock_cursor
            mock_conn.return_value = mock_connection

            orgs = SnowflakeOrgs()
            orgs.wrangles_add("test.sql", "CXO88888", "US", dry_run=True)

            # In dry-run, should NOT execute the CREATE VIEW SQL
            create_view_calls = [
                call
                for call in mock_cursor.execute.call_args_list
                if call[0][0].startswith("CREATE")
            ]
            assert len(create_view_calls) == 0, "Dry-run should not execute CREATE VIEW"

            mock_cursor.close.assert_called()
            mock_connection.close.assert_called()


class TestSnowflakeOrgsWranglesComplete:
    """Comprehensive wrangle tests to reach coverage threshold"""

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt"
    )  # Add this!
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_add_new_view(
        self, mock_exists, mock_conn, mock_profile, mock_token, mock_prompt
    ):
        """Test adding a new wrangle view (no existing grants)"""
        mock_exists.return_value = True
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "yes"  # Auto-confirm the warning

        sql_content = (
            "CREATE OR REPLACE SECURE VIEW CX_DATA_LAKE.CRUX_WRANGLE."
            '"CX09815--TEST_TABLE--DATASET_ID" AS '
            "SELECT * FROM X WHERE \"CDU_ID\"='09815'"
        )

        with patch("builtins.open", mock_open(read_data=sql_content)):
            # Mock SHOW GRANTS returns nothing (view doesn't exist)
            mock_cursor = create_cursor_mock_with_iterations(
                [],  # No existing grants
                [
                    {"name": "CX00000: (CRUX) WRANGLES FOR CXO88888 - CX_DATA_LAKE"}
                ],  # Shares
            )

            mock_connection = MagicMock()
            mock_connection.cursor.return_value = mock_cursor
            mock_conn.return_value = mock_connection

            orgs = SnowflakeOrgs()
            orgs.wrangles_add("test.sql", "CXO88888", "US", dry_run=False)

            # Verify CREATE VIEW was executed
            create_calls = [
                call
                for call in mock_cursor.execute.call_args_list
                if "CREATE" in str(call[0][0])
            ]
            assert len(create_calls) >= 1

            mock_cursor.close.assert_called()
            mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt"
    )  # Add this!
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_add_with_existing_grants(
        self, mock_exists, mock_conn, mock_profile, mock_token, mock_prompt
    ):
        """Test adding wrangle when it already has grants (re-grant scenario)"""
        mock_exists.return_value = True
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "yes"  # Auto-confirm the warning

        sql_content = (
            "CREATE OR REPLACE SECURE VIEW CX_DATA_LAKE.CRUX_WRANGLE."
            '"CX09815--TEST" AS SELECT * FROM X WHERE "CDU_ID"=\'09815\''
        )

        with patch("builtins.open", mock_open(read_data=sql_content)):
            # Mock existing grants
            mock_cursor = create_cursor_mock_with_iterations(
                [  # Existing grants
                    {
                        "granted_to": "SHARE",
                        "grantee_name": "TS91674.CX00000: "
                        "(CRUX) WRANGLES FOR CXO83397 - CX_DATA_LAKE",
                    },
                    {
                        "granted_to": "SHARE",
                        "grantee_name": "CX00000: (CRUX) WRANGLES FOR CXO88888 - CX_DATA_LAKE",
                    },
                ],
                [
                    {"name": "CX00000: (CRUX) WRANGLES FOR CXO88888 - CX_DATA_LAKE"}
                ],  # Shares
            )

            mock_connection = MagicMock()
            mock_connection.cursor.return_value = mock_cursor
            mock_conn.return_value = mock_connection

            orgs = SnowflakeOrgs()
            orgs.wrangles_add("test.sql", "CXO88888", "US", dry_run=False)

            # Should re-grant to CXO83397
            grant_calls = [
                call
                for call in mock_cursor.execute.call_args_list
                if "GRANT" in str(call[0][0])
            ]
            assert len(grant_calls) >= 2  # At least 2 grants

            mock_cursor.close.assert_called()
            mock_connection.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("os.path.exists")
    def test_wrangles_add_cdu_warning(
        self, mock_exists, mock_conn, mock_profile, mock_token
    ):
        """Test warning when CDU_ID has CX prefix in WHERE clause"""
        mock_exists.return_value = True
        mock_token.return_value = "test_token"

        # SQL with CX prefix (should trigger warning)
        sql_content = (
            'CREATE VIEW CX_DATA_LAKE.CRUX_WRANGLE."CX09815--TEST" AS '
            "SELECT * FROM X WHERE \"CDU_ID\"='CX09815'"
        )

        with patch("builtins.open", mock_open(read_data=sql_content)):
            with patch(
                "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt"
            ) as mock_prompt:
                mock_prompt.return_value = "no"  # User says no to continue

                mock_cursor = MagicMock()
                mock_connection = MagicMock()
                mock_connection.cursor.return_value = mock_cursor
                mock_conn.return_value = mock_connection

                orgs = SnowflakeOrgs()

                with patch(
                    "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.Exit"
                ) as mock_exit:
                    mock_exit.side_effect = Exception("Exit called")
                    try:
                        orgs.wrangles_add("test.sql", "CXO88888", "US", dry_run=False)
                    except Exception as e:
                        assert str(e) == "Exit called"

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_wrangles_list_no_views(self, mock_conn):
        """Test listing wrangles when none exist"""
        mock_cursor = create_cursor_mock_with_iterations([])  # No grants = no views

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_list("CXO88888", "US")

        mock_cursor.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_wrangles_list_multiple_views(self, mock_conn):
        """Test listing multiple wrangle views"""
        mock_cursor = create_cursor_mock_with_iterations(
            [  # Multiple views
                {
                    "granted_on": "VIEW",
                    "name": 'CX_DATA_LAKE.CRUX_WRANGLE."CX09815--TEST1"',
                },
                {
                    "granted_on": "VIEW",
                    "name": 'CX_DATA_LAKE.CRUX_WRANGLE."CX09815--TEST2"',
                },
                {
                    "granted_on": "VIEW",
                    "name": 'CX_DATA_LAKE.CRUX_WRANGLE."CX09815--TEST3"',
                },
            ]
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_list("CXO88888", "US")

        mock_cursor.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_wrangles_share_add_success(self, mock_conn, mock_profile, mock_token):
        """Test successfully adding wrangle to share"""
        mock_token.return_value = "test_token"

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_share_add("CX09815--TEST--DATASET", "CXO88888", "US")

        # Should execute GRANT
        grant_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "GRANT" in str(call[0][0])
        ]
        assert len(grant_calls) >= 1

        mock_cursor.close.assert_called()

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.requests.get")
    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.typer.prompt")
    def test_wrangles_share_remove_success(
        self, mock_prompt, mock_requests, mock_conn, mock_profile, mock_token
    ):
        """Test successfully removing wrangle from share"""
        mock_token.return_value = "test_token"
        mock_prompt.return_value = "REVOKE"

        mock_response = MagicMock()
        mock_response.json.return_value = {"roles": ["ciam-role-CruxAdmin"]}
        mock_requests.return_value = mock_response

        mock_cursor = MagicMock()
        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_share_remove("CX09815--TEST--DATASET", "CXO88888", "US")

        # Should execute REVOKE
        revoke_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "REVOKE" in str(call[0][0])
        ]
        assert len(revoke_calls) >= 1

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_wrangles_share_create_new(self, mock_conn, mock_profile, mock_token):
        """Test creating a new wrangle share"""
        mock_token.return_value = "test_token"

        mock_cursor = create_cursor_mock_with_iterations([])  # No existing shares

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_share_create("CXO88888", "KMB96641", "US")

        # Should execute CREATE SHARE and GRANT commands
        create_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "CREATE SHARE" in str(call[0][0])
        ]
        assert len(create_calls) >= 1

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_wrangles_share_create_already_exists(
        self, mock_conn, mock_profile, mock_token
    ):
        """Test creating wrangle share when it already exists"""
        mock_token.return_value = "test_token"

        mock_cursor = create_cursor_mock_with_iterations(
            [
                {"name": "CX00000: (CRUX) WRANGLES FOR CXO88888 - CX_DATA_LAKE"}
            ]  # Already exists
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.wrangles_share_create("CXO88888", "KMB96641", "US")

        # Should NOT create new share
        create_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "CREATE SHARE" in str(call[0][0])
        ]
        assert len(create_calls) == 0

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_entitlements_sync_from_snowflake_success(
        self, mock_conn, mock_profile, mock_token
    ):
        """Test Snowflake-based entitlements sync"""
        mock_token.return_value = "test_token"

        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}, {"ACCOUNT": "ZSB87868"}],  # 2 accounts
            [{"CDU_ID": "03243"}, {"CDU_ID": "03249"}],  # 2 CDUs
            [{"SVC_ACCT": "KMB96641", "CDU_ID": "03243"}],  # 1 exists
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync_from_snowflake("CXO88888", dry_run=False)

        # Should insert 3 new restrictions (2 accounts × 2 CDUs - 1 existing = 3)
        insert_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "INSERT" in str(call[0][0])
            and "CRUX_SERVICE_RESTRICTIONS" in str(call[0][0])
        ]
        assert len(insert_calls) == 3

    @patch("cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.set_api_token")
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_current_profile"
    )
    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs.get_snowflake_connection_instance"
    )
    def test_entitlements_sync_from_snowflake_dry_run(
        self, mock_conn, mock_profile, mock_token
    ):
        """Test Snowflake-based sync in dry-run mode"""
        mock_token.return_value = "test_token"

        mock_cursor = create_cursor_mock_with_iterations(
            [{"ACCOUNT": "KMB96641"}],  # 1 account
            [{"CDU_ID": "03243"}],  # 1 CDU
            [],  # No existing
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.entitlements_sync_from_snowflake("CXO88888", dry_run=True)

        # Should NOT insert anything
        insert_calls = [
            call
            for call in mock_cursor.execute.call_args_list
            if "INSERT" in str(call[0][0])
        ]
        assert len(insert_calls) == 0


class TestSnowflakeOrgsListAll:
    """Tests for list_all method with --unique flag"""

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_list_all_default_with_accounts(self, mock_conn):
        """Test list_all without unique flag (shows accounts)"""
        mock_cursor = create_cursor_mock_with_iterations(
            [
                {"ORG": "CXO01010", "NAME": "Test Org 1", "ACCOUNT": "KMB96641"},
                {"ORG": "CXO01010", "NAME": "Test Org 1", "ACCOUNT": "ZSB87868"},
                {"ORG": "CXO02020", "NAME": "Test Org 2", "ACCOUNT": "ABC12345"},
            ]
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.list_all(unique=False)

        # Verify SQL includes ACCOUNT column
        executed_sql = mock_cursor.execute.call_args[0][0]
        assert '"ACCOUNT"' in executed_sql
        assert 'ORDER BY "NAME", "ACCOUNT" ASC' in executed_sql

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_list_all_unique_without_accounts(self, mock_conn):
        """Test list_all with unique flag (excludes accounts)"""
        mock_cursor = create_cursor_mock_with_iterations(
            [
                {"ORG": "CXO01010", "NAME": "Test Org 1"},
                {"ORG": "CXO02020", "NAME": "Test Org 2"},
                {"ORG": "CXO03030", "NAME": "Test Org 3"},
            ]
        )

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.list_all(unique=True)

        # Verify SQL excludes ACCOUNT column
        executed_sql = mock_cursor.execute.call_args[0][0]
        assert '"ACCOUNT"' not in executed_sql
        assert 'ORDER BY "NAME" ASC' in executed_sql
        assert 'ORDER BY "NAME", "ACCOUNT"' not in executed_sql

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_list_all_no_organizations_found(self, mock_conn):
        """Test list_all when no organizations exist"""
        mock_cursor = create_cursor_mock_with_iterations([])  # Empty result

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.list_all(unique=False)

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()

    @patch(
        "cruxctl.command_groups.bigsnow.snowflake.snowflake_orgs."
        "get_snowflake_connection_instance"
    )
    def test_list_all_unique_no_organizations_found(self, mock_conn):
        """Test list_all with unique flag when no organizations exist"""
        mock_cursor = create_cursor_mock_with_iterations([])  # Empty result

        mock_connection = MagicMock()
        mock_connection.cursor.return_value = mock_cursor
        mock_conn.return_value = mock_connection

        orgs = SnowflakeOrgs()
        orgs.list_all(unique=True)

        mock_cursor.close.assert_called()
        mock_connection.close.assert_called()
