- Add tests for both initialization methods of VLLMClient

- Verify correct behavior of base_url attribute
- Update tests to include tensor parallelism scenarios
- Ensure proper cleanup of client resources

Author: re-imagined

tests/test_vllm_client_server.py

@@ -43,8 +43,32 @@ def setUpClass(cls):
             ["trl", "vllm-serve", "--model", cls.model_id], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
         )
 
-        # Initialize the client
-        cls.client = VLLMClient(connection_timeout=120)
+        # Initialize the clients using both initialization methods
+        cls.client = VLLMClient(connection_timeout=120)  # Default host and port
+        cls.client_base_url = VLLMClient(base_url="http://0.0.0.0:8000", connection_timeout=120)  # Using base_url
+
+    def test_initialization_methods(self):
+        """Test that both initialization methods work correctly."""
+        # Test generation with default client (host+port)
+        prompts = ["Test initialization 1"]
+        outputs_default = self.client.generate(prompts)
+        self.assertIsInstance(outputs_default, list)
+        self.assertEqual(len(outputs_default), len(prompts))
+
+        # Test generation with base_url client
+        outputs_base_url = self.client_base_url.generate(prompts)
+        self.assertIsInstance(outputs_base_url, list)
+        self.assertEqual(len(outputs_base_url), len(prompts))
+
+    def test_base_url_attribute(self):
+        """Test that both initialization methods set the base_url attribute correctly."""
+        # Both clients should have the same base_url
+        self.assertEqual(self.client.base_url, "http://0.0.0.0:8000")
+        self.assertEqual(self.client_base_url.base_url, "http://0.0.0.0:8000")
+
+        # Verify the client doesn't store host/port when base_url is provided
+        self.assertTrue(not hasattr(self.client_base_url, 'host') or self.client_base_url.host is None)
+        self.assertTrue(not hasattr(self.client_base_url, 'server_port') or self.client_base_url.server_port is None)
 
     def test_generate(self):
         prompts = ["Hello, AI!", "Tell me a joke"]
@@ -90,8 +114,9 @@ def test_reset_prefix_cache(self):
     def tearDownClass(cls):
         super().tearDownClass()
 
-        # Close the client
+        # Close the clients
         cls.client.close_communicator()
+        cls.client_base_url.close_communicator()
 
         # vLLM x pytest (or Popen) seems not to handle process termination well. To avoid zombie processes, we need to
         # kill the server process and its children explicitly.
@@ -122,8 +147,32 @@ def setUpClass(cls):
             env=env,
         )
 
-        # Initialize the client
-        cls.client = VLLMClient(connection_timeout=120)
+        # Initialize the clients using both initialization methods
+        cls.client = VLLMClient(connection_timeout=120)  # Default host and port
+        cls.client_base_url = VLLMClient(base_url="http://0.0.0.0:8000", connection_timeout=120)  # Using base_url
+
+    def test_initialization_methods(self):
+        """Test that both initialization methods work correctly with tensor parallelism enabled."""
+        # Test generation with default client (host+port)
+        prompts = ["Test TP initialization 1"]
+        outputs_default = self.client.generate(prompts)
+        self.assertIsInstance(outputs_default, list)
+        self.assertEqual(len(outputs_default), len(prompts))
+
+        # Test generation with base_url client
+        outputs_base_url = self.client_base_url.generate(prompts)
+        self.assertIsInstance(outputs_base_url, list)
+        self.assertEqual(len(outputs_base_url), len(prompts))
+
+    def test_base_url_attribute(self):
+        """Test that both initialization methods set the base_url attribute correctly."""
+        # Both clients should have the same base_url
+        self.assertEqual(self.client.base_url, "http://0.0.0.0:8000")
+        self.assertEqual(self.client_base_url.base_url, "http://0.0.0.0:8000")
+
+        # Verify the client doesn't store host/port when base_url is provided
+        self.assertTrue(not hasattr(self.client_base_url, 'host') or self.client_base_url.host is None)
+        self.assertTrue(not hasattr(self.client_base_url, 'server_port') or self.client_base_url.server_port is None)
 
     def test_generate(self):
         prompts = ["Hello, AI!", "Tell me a joke"]
@@ -151,8 +200,9 @@ def test_reset_prefix_cache(self):
     def tearDownClass(cls):
         super().tearDownClass()
 
-        # Close the client
+        # Close the clients
         cls.client.close_communicator()
+        cls.client_base_url.close_communicator()
 
         # vLLM x pytest (or Popen) seems not to handle process termination well. To avoid zombie processes, we need to
         # kill the server process and its children explicitly.