diff --git a/examples/basic_example.py b/examples/basic_example.py index 0366b6d..7cf62b3 100644 --- a/examples/basic_example.py +++ b/examples/basic_example.py @@ -71,11 +71,23 @@ def main(): ) print(f"✓ Customer created/retrieved - ID: {customer.customer_id}") + # Display service token after customer creation + token_after_customer = client.state_manager.get_dynamic_token() + if token_after_customer: + print(f" Service token: {token_after_customer}") + # Create or get instance print("\nCreating/getting instance for customer...") instance = customer.get_or_create_instance() print(f"✓ Instance created/retrieved - ID: {instance.instance_id}") + # Display service token after instance creation (may have been replaced) + token_after_instance = client.state_manager.get_dynamic_token() + if token_after_instance: + print(f" Service token: {token_after_instance}") + if token_after_customer != token_after_instance: + print(" ⚠️ Token was replaced by instance-specific token") + # Set instance status instance.set_status(args.status) print(f"✓ Instance status set to: {args.status}") @@ -89,6 +101,14 @@ def main(): print(f"Customer ID: {customer.customer_id}") print(f"Instance ID: {instance.instance_id}") + # Show final token + final_token = client.state_manager.get_dynamic_token() + print("\nService Token Information:") + if final_token: + print(f" Active service token: {final_token}") + else: + print(" Service token: Not available") + if __name__ == "__main__": main() diff --git a/examples/metrics_example.py b/examples/metrics_example.py index f0a37bc..5993f0b 100644 --- a/examples/metrics_example.py +++ b/examples/metrics_example.py @@ -72,14 +72,21 @@ async def main(): ) print(f"✓ Customer created/retrieved - ID: {customer.customer_id}") + # Display service token after customer creation + token_after_customer = client.state_manager.get_dynamic_token() + if token_after_customer: + print(f" Service token: {token_after_customer}") + # Get or create the associated instance instance = await customer.get_or_create_instance() - print(f"Instance ID: {instance.instance_id}") print(f"✓ Instance created/retrieved - ID: {instance.instance_id}") - # Get or create the associated instance - instance = await customer.get_or_create_instance() - print(f"Instance ID: {instance.instance_id}") + # Display service token after instance creation (may have been replaced) + token_after_instance = client.state_manager.get_dynamic_token() + if token_after_instance: + print(f" Service token: {token_after_instance}") + if token_after_customer != token_after_instance: + print(" ⚠️ Token was replaced by instance-specific token") # Set instance status await instance.set_status(args.status) @@ -96,9 +103,19 @@ async def main(): instance.send_metric("memory_usage", 0.67), instance.send_metric("disk_usage", 0.45), ) - print("Metrics sent successfully") + print("✓ Metrics sent successfully") + + print("\n🎉 Metrics example completed successfully!") + print(f"Customer ID: {customer.customer_id}") + print(f"Instance ID: {instance.instance_id}") - print(f"Instance ID: {instance.instance_id}") + # Show final token + final_token = client.state_manager.get_dynamic_token() + print("\nService Token Information:") + if final_token: + print(f" Active service token: {final_token}") + else: + print(" Service token: Not available") if __name__ == "__main__": diff --git a/replicated/resources.py b/replicated/resources.py index 925c52c..97176ab 100644 --- a/replicated/resources.py +++ b/replicated/resources.py @@ -27,14 +27,26 @@ def __init__( self.channel = channel self._data = kwargs - def get_or_create_instance(self) -> Union["Instance", "AsyncInstance"]: + def get_or_create_instance( + self, service_account_token: Optional[str] = None + ) -> Union["Instance", "AsyncInstance"]: """Get or create an instance for this customer.""" if hasattr(self._client, "_get_or_create_instance_async"): # type: ignore[arg-type] - return AsyncInstance(self._client, self.customer_id, self.instance_id) + return AsyncInstance( + self._client, + self.customer_id, + self.instance_id, + service_account_token=service_account_token, + ) else: # type: ignore[arg-type] - return Instance(self._client, self.customer_id, self.instance_id) + return Instance( + self._client, + self.customer_id, + self.instance_id, + service_account_token=service_account_token, + ) def __getattr__(self, name: str) -> Any: """Access additional customer data.""" @@ -45,10 +57,17 @@ class AsyncCustomer(Customer): """Async version of Customer.""" # type: ignore[override] - async def get_or_create_instance(self) -> "AsyncInstance": + async def get_or_create_instance( + self, service_account_token: Optional[str] = None + ) -> "AsyncInstance": """Get or create an instance for this customer.""" # type: ignore[arg-type] - return AsyncInstance(self._client, self.customer_id, self.instance_id) + return AsyncInstance( + self._client, + self.customer_id, + self.instance_id, + service_account_token=service_account_token, + ) class Instance: @@ -59,11 +78,13 @@ def __init__( client: "ReplicatedClient", customer_id: str, instance_id: Optional[str] = None, + service_account_token: Optional[str] = None, **kwargs: Any, ) -> None: self._client = client self.customer_id = customer_id self.instance_id = instance_id + self._service_account_token = service_account_token self._machine_id = client._machine_id self._data = kwargs self._status = "ready" @@ -112,12 +133,23 @@ def set_version(self, version: str) -> None: def _ensure_instance(self) -> None: """Ensure the instance ID is generated and cached.""" if self.instance_id: + # If we have an instance ID but a service token was provided, + # replace dynamic token + if self._service_account_token: + self._client.state_manager.set_dynamic_token( + self._service_account_token + ) return # Check if instance ID is cached cached_instance_id = self._client.state_manager.get_instance_id() if cached_instance_id: self.instance_id = cached_instance_id + # If we have a service token provided, replace dynamic token + if self._service_account_token: + self._client.state_manager.set_dynamic_token( + self._service_account_token + ) return # Create new instance @@ -135,6 +167,13 @@ def _ensure_instance(self) -> None: self.instance_id = response["instance_id"] self._client.state_manager.set_instance_id(self.instance_id) + # If API returns a service_token, replace the dynamic token with it + if "service_token" in response: + self._client.state_manager.set_dynamic_token(response["service_token"]) + # Otherwise, if user provided a service token, use that + elif self._service_account_token: + self._client.state_manager.set_dynamic_token(self._service_account_token) + def _report_instance(self) -> None: """Send instance telemetry to vandoor.""" if not self.instance_id: @@ -190,11 +229,13 @@ def __init__( client: "AsyncReplicatedClient", customer_id: str, instance_id: Optional[str] = None, + service_account_token: Optional[str] = None, **kwargs: Any, ) -> None: self._client = client self.customer_id = customer_id self.instance_id = instance_id + self._service_account_token = service_account_token self._machine_id = client._machine_id self._data = kwargs self._status = "ready" @@ -243,12 +284,23 @@ async def set_version(self, version: str) -> None: async def _ensure_instance(self) -> None: """Ensure the instance ID is generated and cached.""" if self.instance_id: + # If we have an instance ID but a service token was provided, + # replace dynamic token + if self._service_account_token: + self._client.state_manager.set_dynamic_token( + self._service_account_token + ) return # Check if instance ID is cached cached_instance_id = self._client.state_manager.get_instance_id() if cached_instance_id: self.instance_id = cached_instance_id + # If we have a service token provided, replace dynamic token + if self._service_account_token: + self._client.state_manager.set_dynamic_token( + self._service_account_token + ) return # Create new instance @@ -266,6 +318,13 @@ async def _ensure_instance(self) -> None: self.instance_id = response["instance_id"] self._client.state_manager.set_instance_id(self.instance_id) + # If API returns a service_token, replace the dynamic token with it + if "service_token" in response: + self._client.state_manager.set_dynamic_token(response["service_token"]) + # Otherwise, if user provided a service token, use that + elif self._service_account_token: + self._client.state_manager.set_dynamic_token(self._service_account_token) + async def _report_instance(self) -> None: """Send instance telemetry to vandoor.""" if not self.instance_id: diff --git a/tests/test_client.py b/tests/test_client.py index fbdecab..e23fb45 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -160,6 +160,141 @@ def test_instance_uses_machine_id_in_headers(self, mock_httpx): assert "X-Replicated-ClusterID" in headers assert headers["X-Replicated-ClusterID"] == client._machine_id + def test_instance_with_service_account_token(self): + """Test that instances can be created with a service account token.""" + from replicated.resources import Instance + + client = ReplicatedClient(publishable_key="pk_test_123", app_slug="my-app") + instance = Instance( + client, + "customer_123", + "instance_123", + service_account_token="test_token_123", + ) + + assert instance._service_account_token == "test_token_123" + + @patch("replicated.http_client.httpx.Client") + def test_get_or_create_instance_with_service_token(self, mock_httpx): + """Test that get_or_create_instance passes service token to instance.""" + mock_response = Mock() + mock_response.is_success = True + mock_response.json.return_value = { + "customer": { + "id": "customer_123", + "email": "test@example.com", + "name": "test user", + "serviceToken": "service_token_123", + "instanceId": "instance_123", + } + } + + mock_client = Mock() + mock_client.request.return_value = mock_response + mock_httpx.return_value = mock_client + + client = ReplicatedClient(publishable_key="pk_test_123", app_slug="my-app") + customer = client.customer.get_or_create("test@example.com") + instance = customer.get_or_create_instance( + service_account_token="instance_token_abc" + ) + + assert instance._service_account_token == "instance_token_abc" + + @patch("replicated.http_client.httpx.Client") + def test_ensure_instance_replaces_dynamic_token_from_api(self, mock_httpx): + """Test that _ensure_instance replaces dynamic_token with + service_token from API.""" + from replicated.resources import Instance + + with tempfile.TemporaryDirectory() as tmpdir: + mock_response = Mock() + mock_response.is_success = True + mock_response.json.return_value = { + "instance_id": "instance_789", + "service_token": "api_returned_token_xyz", + } + + mock_client = Mock() + mock_client.request.return_value = mock_response + mock_httpx.return_value = mock_client + + client = ReplicatedClient( + publishable_key="pk_test_123", + app_slug="my-app", + state_directory=tmpdir, + ) + + # Set initial customer token + client.state_manager.set_dynamic_token("customer_token_abc") + + instance = Instance(client, "customer_123") + + # Trigger instance creation + instance._ensure_instance() + + # Verify dynamic token was replaced with instance token + stored_token = client.state_manager.get_dynamic_token() + assert stored_token == "api_returned_token_xyz" + + @patch("replicated.http_client.httpx.Client") + def test_service_account_token_replaces_dynamic_token(self, mock_httpx): + """Test that providing service_account_token replaces the dynamic_token.""" + from replicated.resources import Instance + + with tempfile.TemporaryDirectory() as tmpdir: + mock_response = Mock() + mock_response.is_success = True + mock_response.json.return_value = { + "instance_id": "instance_789", + } + + mock_client = Mock() + mock_client.request.return_value = mock_response + mock_httpx.return_value = mock_client + + client = ReplicatedClient( + publishable_key="pk_test_123", + app_slug="my-app", + state_directory=tmpdir, + ) + + # Set initial customer token + client.state_manager.set_dynamic_token("customer_token_abc") + + instance = Instance( + client, "customer_123", service_account_token="user_provided_token" + ) + + # Trigger instance creation + instance._ensure_instance() + + # Verify dynamic token was replaced with user-provided token + stored_token = client.state_manager.get_dynamic_token() + assert stored_token == "user_provided_token" + + @patch("replicated.http_client.httpx.Client") + def test_auth_headers_use_dynamic_token(self, mock_httpx): + """Test that _get_auth_headers uses the dynamic token.""" + with tempfile.TemporaryDirectory() as tmpdir: + mock_client = Mock() + mock_httpx.return_value = mock_client + + client = ReplicatedClient( + publishable_key="pk_test_123", + app_slug="my-app", + state_directory=tmpdir, + ) + + # No token set, should use publishable key + headers = client._get_auth_headers() + assert headers["Authorization"] == "Bearer pk_test_123" + + # Set dynamic token (could be from customer or instance) + client.state_manager.set_dynamic_token("dynamic_token_xyz") + headers = client._get_auth_headers() + assert headers["Authorization"] == "dynamic_token_xyz" + class TestAsyncReplicatedClient: @pytest.mark.asyncio @@ -301,3 +436,50 @@ async def test_instance_uses_machine_id_in_headers(self): headers = call_args[1]["headers"] assert "X-Replicated-ClusterID" in headers assert headers["X-Replicated-ClusterID"] == client._machine_id + + @pytest.mark.asyncio + async def test_instance_with_service_account_token(self): + """Test that async instances can be created with a service account token.""" + from replicated.resources import AsyncInstance + + client = AsyncReplicatedClient(publishable_key="pk_test_123", app_slug="my-app") + instance = AsyncInstance( + client, + "customer_123", + "instance_123", + service_account_token="test_token_123", + ) + + assert instance._service_account_token == "test_token_123" + + @pytest.mark.asyncio + async def test_get_or_create_instance_with_service_token(self): + """Test that async get_or_create_instance passes service token to instance.""" + with patch("replicated.http_client.httpx.AsyncClient") as mock_httpx: + from unittest.mock import AsyncMock + + mock_response = Mock() + mock_response.is_success = True + mock_response.json.return_value = { + "customer": { + "id": "customer_123", + "email": "test@example.com", + "name": "test user", + "serviceToken": "service_token_123", + "instanceId": "instance_123", + } + } + + mock_client = Mock() + mock_client.request = AsyncMock(return_value=mock_response) + mock_httpx.return_value = mock_client + + client = AsyncReplicatedClient( + publishable_key="pk_test_123", app_slug="my-app" + ) + customer = await client.customer.get_or_create("test@example.com") + instance = await customer.get_or_create_instance( + service_account_token="instance_token_abc" + ) + + assert instance._service_account_token == "instance_token_abc"