Add tests for BaseSalesforceOrg model and its subclasses

* Test field validation and enum usage
* Add tests for `exchange_token` function using `BaseSalesforceOrg` model
* Test simple-salesforce integration for creating scratch orgs and exchanging auth codes
* Add tests for CLI commands related to creating scratch orgs and exchanging auth codes
* Test command execution and output

tests/test_auth_url.py

@@ -2,8 +2,8 @@
 import unittest
 from unittest.mock import patch, MagicMock
 from pydantic import SecretStr
-from d2x.auth.sf.auth_url import exchange_token
-from d2x.models.sf.org import SalesforceOrgInfo
+from d2x.auth.sf.auth_url import exchange_token, create_scratch_org
+from d2x.models.sf.org import SalesforceOrgInfo, ScratchOrg
 from d2x.base.types import CLIOptions
 from d2x.models.sf.auth import AuthInfo
 
@@ -20,9 +20,15 @@ def test_exchange_token_success(self, mock_https_connection, mock_set_env_var):
                 refresh_token="test_refresh_token",
                 instance_url="https://test.salesforce.com",
             ),
-            org_type="production",
-            domain_type="pod",
-            full_domain="test.salesforce.com",
+            org=ScratchOrg(
+                org_type="scratch",
+                domain_type="my",
+                full_domain="test.salesforce.com",
+                instance_url="https://test.salesforce.com",
+                created_date="2021-01-01",
+                last_modified_date="2021-01-02",
+                status="Active",
+            ),
         )
 
         # Mock the CLIOptions
@@ -69,9 +75,15 @@ def test_exchange_token_failure(self, mock_https_connection, mock_set_env_var):
                 refresh_token="test_refresh_token",
                 instance_url="https://test.salesforce.com",
             ),
-            org_type="production",
-            domain_type="pod",
-            full_domain="test.salesforce.com",
+            org=ScratchOrg(
+                org_type="scratch",
+                domain_type="my",
+                full_domain="test.salesforce.com",
+                instance_url="https://test.salesforce.com",
+                created_date="2021-01-01",
+                last_modified_date="2021-01-02",
+                status="Active",
+            ),
         )
 
         # Mock the CLIOptions
@@ -93,5 +105,77 @@ def test_exchange_token_failure(self, mock_https_connection, mock_set_env_var):
             exchange_token(org_info, cli_options)
 
 
+class TestCreateScratchOrg(unittest.TestCase):
+    @patch("d2x.auth.sf.auth_url.Salesforce")
+    def test_create_scratch_org_success(self, mock_salesforce):
+        # Mock the ScratchOrg
+        org_info = ScratchOrg(
+            org_type="scratch",
+            domain_type="my",
+            full_domain="test.salesforce.com",
+            instance_url="https://test.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+            auth_info=AuthInfo(
+                client_id="test_client_id",
+                client_secret=SecretStr("test_client_secret"),
+                refresh_token="test_refresh_token",
+                instance_url="https://test.salesforce.com",
+            ),
+        )
+
+        # Mock the CLIOptions
+        cli_options = CLIOptions(output_format="text", debug=False)
+
+        # Mock the Salesforce instance and response
+        mock_sf_instance = MagicMock()
+        mock_salesforce.return_value = mock_sf_instance
+        mock_sf_instance.restful.return_value = {
+            "id": "test_org_id",
+            "status": "Active",
+            "expirationDate": "2021-01-10",
+        }
+
+        # Call the function
+        result = create_scratch_org(org_info, cli_options)
+
+        # Assertions
+        self.assertEqual(result["id"], "test_org_id")
+        self.assertEqual(result["status"], "Active")
+        self.assertEqual(result["expirationDate"], "2021-01-10")
+
+    @patch("d2x.auth.sf.auth_url.Salesforce")
+    def test_create_scratch_org_failure(self, mock_salesforce):
+        # Mock the ScratchOrg
+        org_info = ScratchOrg(
+            org_type="scratch",
+            domain_type="my",
+            full_domain="test.salesforce.com",
+            instance_url="https://test.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+            auth_info=AuthInfo(
+                client_id="test_client_id",
+                client_secret=SecretStr("test_client_secret"),
+                refresh_token="test_refresh_token",
+                instance_url="https://test.salesforce.com",
+            ),
+        )
+
+        # Mock the CLIOptions
+        cli_options = CLIOptions(output_format="text", debug=False)
+
+        # Mock the Salesforce instance and response
+        mock_sf_instance = MagicMock()
+        mock_salesforce.return_value = mock_sf_instance
+        mock_sf_instance.restful.side_effect = Exception("Failed to create scratch org")
+
+        # Call the function and assert exception
+        with self.assertRaises(Exception):
+            create_scratch_org(org_info, cli_options)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_cli.py

@@ -0,0 +1,29 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from click.testing import CliRunner
+from d2x.cli.main import d2x_cli
+from d2x.models.sf.org import ScratchOrg
+from d2x.models.sf.auth import AuthInfo
+from d2x.base.types import CLIOptions
+
+
+class TestCLICommands(unittest.TestCase):
+    @patch("d2x.cli.main.create_scratch_org")
+    def test_create_scratch_org_command(self, mock_create_scratch_org):
+        runner = CliRunner()
+        result = runner.invoke(d2x_cli, ["sf", "org", "create"])
+
+        self.assertEqual(result.exit_code, 0)
+        mock_create_scratch_org.assert_called_once()
+
+    @patch("d2x.cli.main.exchange_token")
+    def test_exchange_token_command(self, mock_exchange_token):
+        runner = CliRunner()
+        result = runner.invoke(d2x_cli, ["sf", "org", "exchange"])
+
+        self.assertEqual(result.exit_code, 0)
+        mock_exchange_token.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_org_models.py

@@ -0,0 +1,101 @@
+import unittest
+from pydantic import ValidationError
+from d2x.models.sf.org import (
+    BaseSalesforceOrg,
+    ProductionOrg,
+    TrialOrg,
+    SandboxOrg,
+    ScratchOrg,
+    OrgType,
+    DomainType,
+)
+
+
+class TestBaseSalesforceOrg(unittest.TestCase):
+    def test_production_org(self):
+        org = ProductionOrg(
+            org_type=OrgType.PRODUCTION,
+            domain_type=DomainType.POD,
+            full_domain="example.salesforce.com",
+            instance_url="https://example.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+        )
+        self.assertEqual(org.org_type, OrgType.PRODUCTION)
+        self.assertEqual(org.domain_type, DomainType.POD)
+        self.assertEqual(org.full_domain, "example.salesforce.com")
+        self.assertEqual(org.instance_url, "https://example.salesforce.com")
+        self.assertEqual(org.created_date, "2021-01-01")
+        self.assertEqual(org.last_modified_date, "2021-01-02")
+        self.assertEqual(org.status, "Active")
+
+    def test_trial_org(self):
+        org = TrialOrg(
+            org_type=OrgType.DEMO,
+            domain_type=DomainType.POD,
+            full_domain="example.salesforce.com",
+            instance_url="https://example.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+        )
+        self.assertEqual(org.org_type, OrgType.DEMO)
+        self.assertEqual(org.domain_type, DomainType.POD)
+        self.assertEqual(org.full_domain, "example.salesforce.com")
+        self.assertEqual(org.instance_url, "https://example.salesforce.com")
+        self.assertEqual(org.created_date, "2021-01-01")
+        self.assertEqual(org.last_modified_date, "2021-01-02")
+        self.assertEqual(org.status, "Active")
+
+    def test_sandbox_org(self):
+        org = SandboxOrg(
+            org_type=OrgType.SANDBOX,
+            domain_type=DomainType.POD,
+            full_domain="example.salesforce.com",
+            instance_url="https://example.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+        )
+        self.assertEqual(org.org_type, OrgType.SANDBOX)
+        self.assertEqual(org.domain_type, DomainType.POD)
+        self.assertEqual(org.full_domain, "example.salesforce.com")
+        self.assertEqual(org.instance_url, "https://example.salesforce.com")
+        self.assertEqual(org.created_date, "2021-01-01")
+        self.assertEqual(org.last_modified_date, "2021-01-02")
+        self.assertEqual(org.status, "Active")
+
+    def test_scratch_org(self):
+        org = ScratchOrg(
+            org_type=OrgType.SCRATCH,
+            domain_type=DomainType.POD,
+            full_domain="example.salesforce.com",
+            instance_url="https://example.salesforce.com",
+            created_date="2021-01-01",
+            last_modified_date="2021-01-02",
+            status="Active",
+        )
+        self.assertEqual(org.org_type, OrgType.SCRATCH)
+        self.assertEqual(org.domain_type, DomainType.POD)
+        self.assertEqual(org.full_domain, "example.salesforce.com")
+        self.assertEqual(org.instance_url, "https://example.salesforce.com")
+        self.assertEqual(org.created_date, "2021-01-01")
+        self.assertEqual(org.last_modified_date, "2021-01-02")
+        self.assertEqual(org.status, "Active")
+
+    def test_invalid_org_type(self):
+        with self.assertRaises(ValidationError):
+            ProductionOrg(
+                org_type="invalid_org_type",
+                domain_type=DomainType.POD,
+                full_domain="example.salesforce.com",
+                instance_url="https://example.salesforce.com",
+                created_date="2021-01-01",
+                last_modified_date="2021-01-02",
+                status="Active",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()