test: improve lazy client test to follow project conventions

- Remove verbose comments and print statements
- Focus on observable behavior rather than internal implementation
- Use proper mocking that matches actual cog integration
- Test that cog.current_scope() is called on client creation
- Address code review feedback from PR discussion

Author: zeke

tests/test_simple_lazy.py

@@ -1,75 +1,52 @@
-"""Simple test showing the lazy client fix works."""
+"""Test lazy client creation in replicate.use()."""
 
 import os
 import sys
-from typing import Any
 from unittest.mock import MagicMock, patch
 
 
-def test_use_does_not_create_client_immediately():
-    """Test that replicate.use() does not create a client until the model is called."""
+def test_use_does_not_raise_without_token():
+    """Test that replicate.use() works even when no API token is available."""
     sys.path.insert(0, "src")
-
-    # Clear any existing token to simulate the original error condition
+    
     with patch.dict(os.environ, {}, clear=True):
         with patch.dict(sys.modules, {"cog": None}):
-            try:
-                import replicate
-
-                # This should work now - no client is created yet
-                model: Any = replicate.use("test/model")  # type: ignore[misc]
-
-                # Verify we got a Function object back
-                from replicate.lib._predictions_use import Function
-
-                assert isinstance(model, Function)
-                print("✓ replicate.use() works without immediate client creation")
-
-                # Verify the client property is a property that will create client on demand
-                # We can't call it without a token, but we can check it's the right type
-                assert hasattr(model, "_client_class")  # type: ignore[misc]
-                print("✓ Client class is stored for lazy creation")
-
-            except Exception as e:
-                print(f"✗ Test failed: {e}")
-                raise
+            import replicate
+            
+            # Should not raise an exception
+            model = replicate.use("test/model")
+            assert model is not None
 
 
-def test_client_created_when_model_called():
-    """Test that the client is created when the model is called."""
+def test_cog_current_scope():
+    """Test that cog.current_scope().context is read on each client creation."""
     sys.path.insert(0, "src")
-
-    # Test that we can create a model function with a token available
-    # Mock cog to provide a token
+    
+    mock_context = MagicMock()
+    mock_context.items.return_value = [("REPLICATE_API_TOKEN", "test-token-1")]
+    
     mock_scope = MagicMock()
-    mock_scope.context = {"REPLICATE_API_TOKEN": "test-token"}
+    mock_scope.context = mock_context
+
     mock_cog = MagicMock()
     mock_cog.current_scope.return_value = mock_scope
 
     with patch.dict(os.environ, {}, clear=True):
         with patch.dict(sys.modules, {"cog": mock_cog}):
             import replicate
-
-            # Create model function - should work without errors
-            model: Any = replicate.use("test/model")  # type: ignore[misc]
-            print("✓ Model function created successfully")
-
-            # Verify the model has the lazy client setup
-            assert hasattr(model, "_client_class")
-            assert isinstance(model._client_class, type)
-            print("✓ Lazy client class is properly configured")
-
-            # Test that accessing _client property works (creates client)
-            try:
-                client = model._client  # This should create the client
-                assert client is not None
-                print("✓ Client created successfully when accessed")
-            except Exception as e:
-                print(f"ℹ  Client creation expected to work but got: {e}")
-                # This is okay - the important thing is that use() worked
-
-
-if __name__ == "__main__":
-    test_use_does_not_create_client_immediately()
-    test_client_created_when_model_called()
-    print("\n✓ All tests passed! The lazy client fix works correctly.")
+            
+            model = replicate.use("test/model")
+            
+            # Access the client property - this should trigger client creation and cog.current_scope call
+            _ = model._client
+
+            assert mock_cog.current_scope.call_count == 1
+
+            # Change the token and access client again - should trigger another call
+            mock_context.items.return_value = [("REPLICATE_API_TOKEN", "test-token-2")]
+            
+            # Create a new model to trigger another client creation
+            model2 = replicate.use("test/model2")
+            _ = model2._client
+                
+            assert mock_cog.current_scope.call_count == 2