fix: add tests

runpod · Oct 29, 2023 · 4a99b23 · 4a99b23
1 parent 329c29e
commit 4a99b23
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 10 deletions.
diff --git a/runpod/cli/groups/project/commands.py b/runpod/cli/groups/project/commands.py
@@ -94,6 +94,10 @@ def deploy_project():
     """ Deploy the project to RunPod. """
     click.echo("Deploying project...")
 
-    create_project_endpoint()
+    endpoint_id = create_project_endpoint()
 
-    click.echo("Project deployed successfully!")
+    click.echo(f"Project deployed successfully! Endpoint ID: {endpoint_id}")
+    click.echo("The following urls are available:")
+    click.echo(f"    - https://api.runpod.ai/v2/{endpoint_id}/runsync")
+    click.echo(f"    - https://api.runpod.ai/v2/{endpoint_id}/run")
+    click.echo(f"    - https://api.runpod.ai/v2/{endpoint_id}/health")
diff --git a/runpod/cli/groups/project/functions.py b/runpod/cli/groups/project/functions.py
@@ -296,19 +296,17 @@ def create_project_endpoint():
         environment_variables[variable] = config['project']['env_vars'][variable]
 
     project_endpoint_template = create_template(
-        name =  f'{config["project"]["name"]}-endpoint ({config["project"]["uuid"]})',
+        name =  f'{config["project"]["name"]}-endpoint | {config["project"]["uuid"]}',
         image_name = config['project']['base_image'],
         container_disk_in_gb = config['project']['container_disk_size_gb'],
-        docker_start_cmd = f'source /runpod-volume/{config["project"]["uuid"]}/venv/bin/activate && python -u /runpod-volume/{config["project"]["uuid"]}/{config["project"]["name"]}/{config["runtime"]["handler_path"]}', # pylint: disable=line-too-long
+        docker_start_cmd = f'bash -c ". /runpod-volume/{config["project"]["uuid"]}/venv/bin/activate && python -u /runpod-volume/{config["project"]["uuid"]}/{config["project"]["name"]}/{config["runtime"]["handler_path"]}"', # pylint: disable=line-too-long
         env = environment_variables, is_serverless = True
     )
 
-    print(project_endpoint_template)
-
     deployed_endpoint = create_endpoint(
-        name = f'{config["project"]["name"]}-endpoint ({config["project"]["uuid"]})',
+        name = f'{config["project"]["name"]}-endpoint | {config["project"]["uuid"]}',
         template_id = project_endpoint_template['id'],
         network_volume_id=config['project']['storage_id'],
     )
 
-    print(deployed_endpoint)
+    return deployed_endpoint['id']
diff --git a/tests/test_cli/test_cli_groups/test_project_commands.py b/tests/test_cli/test_cli_groups/test_project_commands.py
@@ -6,7 +6,9 @@
 
 from click.testing import CliRunner
 from runpod.cli.groups.project.commands import (
-    new_project_wizard, launch_project_pod, start_project_pod)
+    new_project_wizard, launch_project_pod, start_project_pod,
+    deploy_project
+)
 
 class TestProjectCLI(unittest.TestCase):
     ''' A collection of tests for the Project CLI commands. '''
@@ -83,3 +85,15 @@ def test_start_project_pod(self):
 
         self.assertEqual(result.exit_code, 0)
         self.assertIn("Starting project API server...", result.output)
+
+    def test_deploy_project(self):
+        '''
+        Tests the deploy_project command.
+        '''
+        with patch('runpod.cli.groups.project.commands.deploy_project') as mock_deploy:
+            mock_deploy.return_value = None
+            result = self.runner.invoke(deploy_project)
+
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("Deploying project...", result.output)
+        mock_deploy.assert_called_once()
diff --git a/tests/test_cli/test_cli_groups/test_project_functions.py b/tests/test_cli/test_cli_groups/test_project_functions.py
@@ -8,7 +8,8 @@
 
 from runpod.cli.groups.project.functions import(
     STARTER_TEMPLATES, create_new_project,
-    launch_project, start_project_api
+    launch_project, start_project_api,
+    create_project_endpoint
 )
 
 class TestCreateNewProject(unittest.TestCase):
@@ -229,3 +230,45 @@ def test_start_project_api_pod_not_found(self, mock_ssh_connection, mock_get_pro
         )
         assert mock_ssh_connection.called is False
         assert mock_get_project_pod.called
+
+class TestCreateProjectEndpoint(unittest.TestCase):
+    """ Test the create_project_endpoint function. """
+
+    @patch('runpod.cli.groups.project.functions.load_project_config')
+    @patch('runpod.cli.groups.project.functions.create_template')
+    @patch('runpod.cli.groups.project.functions.create_endpoint')
+    def test_create_project_endpoint(self, mock_create_endpoint,
+                                     mock_create_template, mock_load_project_config):
+        """ Test that a project endpoint is created successfully. """
+        mock_load_project_config.return_value = {
+            'project': {
+                'name': 'test_project',
+                'uuid': '123456',
+                'env_vars': {'TEST_VAR': 'value'},
+                'base_image': 'test_image',
+                'container_disk_size_gb': 10,
+                'storage_id': 'test_storage_id',
+            },
+            'runtime': {
+                'handler_path': 'handler.py'
+            }
+        }
+        mock_create_template.return_value = {'id': 'test_template_id'}
+        mock_create_endpoint.return_value = {'id': 'test_endpoint_id'}
+
+        result = create_project_endpoint()
+
+        self.assertEqual(result, 'test_endpoint_id')
+        mock_create_template.assert_called_once_with(
+            name='test_project-endpoint | 123456',
+            image_name='test_image',
+            container_disk_in_gb=10,
+            docker_start_cmd='bash -c ". /runpod-volume/123456/venv/bin/activate && python -u /runpod-volume/123456/test_project/handler.py"', # pylint: disable=line-too-long
+            env={'TEST_VAR': 'value'},
+            is_serverless=True
+        )
+        mock_create_endpoint.assert_called_once_with(
+            name='test_project-endpoint | 123456',
+            template_id='test_template_id',
+            network_volume_id='test_storage_id'
+        )