fix tests

Author: ppandit-sfdc

tests/test_credentials.py

@@ -18,6 +18,7 @@ def test_from_env(self):
             "client_id": "test_client_id",
             "client_secret": "test_secret",
             "login_url": "https://test.login.url",
+            "dataspace": "test_dataspace",  # Added dataspace
         }
 
         with patch.dict(
@@ -30,6 +31,30 @@ def test_from_env(self):
             assert creds.client_id == test_creds["client_id"]
             assert creds.client_secret == test_creds["client_secret"]
             assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace == test_creds["dataspace"]  # Added dataspace assertion
+
+    def test_from_env_without_dataspace(self):
+        """Test loading credentials from environment variables without dataspace."""
+        test_creds = {
+            "username": "test_user",
+            "password": "test_pass",
+            "client_id": "test_client_id",
+            "client_secret": "test_secret",
+            "login_url": "https://test.login.url",
+        }
+
+        # Create env dict without dataspace
+        env_dict = {v: test_creds[k] for k, v in ENV_CREDENTIALS.items() if k != "dataspace"}
+        
+        with patch.dict(os.environ, env_dict, clear=True):
+            creds = Credentials.from_env()
+
+            assert creds.username == test_creds["username"]
+            assert creds.password == test_creds["password"]
+            assert creds.client_id == test_creds["client_id"]
+            assert creds.client_secret == test_creds["client_secret"]
+            assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace is None  # Should default to None
 
     def test_from_env_missing_vars(self):
         """Test that missing environment variables raise appropriate error."""
@@ -47,13 +72,15 @@ def test_from_ini(self):
         client_id = ini_client_id
         client_secret = ini_secret
         login_url = https://ini.login.url
+        dataspace = ini_dataspace
 
         [other_profile]
         username = other_user
         password = other_pass
         client_id = other_client_id
         client_secret = other_secret
         login_url = https://other.login.url
+        dataspace = other_dataspace
         """
 
         with (
@@ -73,6 +100,7 @@ def test_from_ini(self):
                 assert creds.client_id == "ini_client_id"
                 assert creds.client_secret == "ini_secret"
                 assert creds.login_url == "https://ini.login.url"
+                assert creds.dataspace == "ini_dataspace"
 
                 # Test other profile
                 creds = Credentials.from_ini(
@@ -83,6 +111,7 @@ def test_from_ini(self):
                 assert creds.client_id == "other_client_id"
                 assert creds.client_secret == "other_secret"
                 assert creds.login_url == "https://other.login.url"
+                assert creds.dataspace == "other_dataspace"
 
     def test_from_available_env(self):
         """Test that from_available uses environment variables when available."""
@@ -92,6 +121,7 @@ def test_from_available_env(self):
             "client_id": "test_client_id",
             "client_secret": "test_secret",
             "login_url": "https://test.login.url",
+            "dataspace": "test_dataspace",  # Added dataspace
         }
 
         with (
@@ -107,6 +137,7 @@ def test_from_available_env(self):
             assert creds.client_id == test_creds["client_id"]
             assert creds.client_secret == test_creds["client_secret"]
             assert creds.login_url == test_creds["login_url"]
+            assert creds.dataspace == test_creds["dataspace"]  # Added dataspace assertion
 
     def test_from_available_ini(self):
         """Test that from_available uses INI file when env vars not available."""
@@ -117,6 +148,7 @@ def test_from_available_ini(self):
         client_id = ini_client_id
         client_secret = ini_secret
         login_url = https://ini.login.url
+        dataspace = ini_dataspace
         """
 
         with (
@@ -137,6 +169,7 @@ def test_from_available_ini(self):
                 assert creds.client_id == "ini_client_id"
                 assert creds.client_secret == "ini_secret"
                 assert creds.login_url == "https://ini.login.url"
+                assert creds.dataspace == "ini_dataspace"
 
     def test_from_available_no_creds(self):
         """Test that from_available raises error when no credentials are found."""

tests/test_credentials_profile_integration.py

@@ -32,6 +32,7 @@ def test_query_api_reader_with_custom_profile(self):
             mock_credentials.password = "custom_password"
             mock_credentials.client_id = "custom_client_id"
             mock_credentials.client_secret = "custom_secret"
+            mock_credentials.dataspace = "custom_dataspace"  # Added dataspace
             mock_from_available.return_value = mock_credentials
 
             # Mock the SalesforceCDPConnection
@@ -49,13 +50,56 @@ def test_query_api_reader_with_custom_profile(self):
                 # Verify the correct profile was used
                 mock_from_available.assert_called_with(profile="custom_profile")
 
-                # Verify the connection was created with the custom credentials
+                # Verify the connection was created with the custom credentials including dataspace
                 mock_conn_class.assert_called_once_with(
                     "https://custom.salesforce.com",
                     "custom@example.com",
                     "custom_password",
                     "custom_client_id",
                     "custom_secret",
+                    dataspace="custom_dataspace",  # Added dataspace parameter
+                )
+
+    def test_query_api_reader_without_dataspace(self):
+        """Test QueryAPIDataCloudReader works when dataspace is None."""
+        mock_spark = MagicMock()
+
+        with patch(
+            "datacustomcode.credentials.Credentials.from_available"
+        ) as mock_from_available:
+            # Mock credentials without dataspace
+            mock_credentials = MagicMock()
+            mock_credentials.login_url = "https://custom.salesforce.com"
+            mock_credentials.username = "custom@example.com"
+            mock_credentials.password = "custom_password"
+            mock_credentials.client_id = "custom_client_id"
+            mock_credentials.client_secret = "custom_secret"
+            mock_credentials.dataspace = None  # No dataspace
+            mock_from_available.return_value = mock_credentials
+
+            # Mock the SalesforceCDPConnection
+            with patch(
+                "datacustomcode.io.reader.query_api.SalesforceCDPConnection"
+            ) as mock_conn_class:
+                mock_conn = MagicMock()
+                mock_conn_class.return_value = mock_conn
+
+                # Test with custom profile
+                QueryAPIDataCloudReader(
+                    mock_spark, credentials_profile="custom_profile"
+                )
+
+                # Verify the correct profile was used
+                mock_from_available.assert_called_with(profile="custom_profile")
+
+                # Verify the connection was created WITHOUT dataspace parameter when None
+                mock_conn_class.assert_called_once_with(
+                    "https://custom.salesforce.com",
+                    "custom@example.com",
+                    "custom_password",
+                    "custom_client_id",
+                    "custom_secret",
+                    # No dataspace parameter when None
                 )
 
     def test_print_writer_with_custom_profile(self):