Updated image upload tests (#213)

* Updated image upload tests

Signed-off-by: root <root@ip-172-31-53-203.ec2.internal>

* Added back Hugging Face onboarding

Signed-off-by: Michael Lam <michael.lam@aixplain.com>

---------

Signed-off-by: root <root@ip-172-31-53-203.ec2.internal>
Signed-off-by: Michael Lam <michael.lam@aixplain.com>
Co-authored-by: root <root@ip-172-31-53-203.ec2.internal>

aixplain/factories/model_factory.py

@@ -361,7 +361,7 @@ def create_asset_repo(
             if function_dict["name"] == function:
                 function_id = function_dict["id"]
         if function_id is None:
-            raise Exception("Invalid function name")
+            raise Exception(f"Invalid function name {function}")
         create_url = urljoin(config.BACKEND_URL, f"sdk/models/onboard")
         logging.debug(f"URL: {create_url}")
         if api_key:

tests/functional/model/hf_onboarding_test.py

@@ -13,7 +13,7 @@ def test_deploy_model():
     # Start the deployment
     model_name = "Test Model"
     repo_id = "tiiuae/falcon-7b"
-    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN)
+    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN)
     assert "id" in response.keys()
 
     # Check for status
@@ -30,31 +30,31 @@ def test_deploy_model():
     delete_asset(model_id, config.TEAM_API_KEY)
 
 
-@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
+# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
 def test_nonexistent_model():
     # Start the deployment
     model_name = "Test Model"
     repo_id = "nonexistent-supplier/nonexistent-model"
-    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN)
+    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN)
     assert response["statusCode"] == 400
     assert response["message"] == "err.unable_to_onboard_model"
 
 
-@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
+# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
 def test_size_limit():
     # Start the deployment
     model_name = "Test Model"
     repo_id = "tiiuae/falcon-40b"
-    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, config.HF_TOKEN)
+    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token=config.HF_TOKEN)
     assert response["statusCode"] == 400
     assert response["message"] == "err.unable_to_onboard_model"
 
 
-@pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
+# @pytest.mark.skip(reason="Model Deployment is deactivated for improvements.")
 def test_gated_model():
     # Start the deployment
     model_name = "Test Model"
     repo_id = "meta-llama/Llama-2-7b-hf"
-    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, "mock_key")
+    response = ModelFactory.deploy_huggingface_model(model_name, repo_id, hf_token="mock_key")
     assert response["statusCode"] == 400
     assert response["message"] == "err.unable_to_onboard_model"

tests/image_upload_e2e_test.py → tests/functional/model/image_upload_e2e_test.py

@@ -9,7 +9,6 @@
 import pytest
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_create_and_upload_model():
     # List the host machines
     host_response = ModelFactory.list_host_machines()
@@ -31,14 +30,15 @@ def test_create_and_upload_model():
 
     # Register the model, and create an image repository for it.
     with open(Path("tests/test_requests/create_asset_request.json")) as f:
-        register_payload = json.load(f)
-    name = register_payload["name"]
-    host_machine = register_payload["hostingMachine"]
-    version = register_payload["version"]
-    description = register_payload["description"]
-    function = register_payload["function"]
-    source_language = register_payload["sourceLanguage"]
-    register_response = ModelFactory.create_asset_repo(name, host_machine, version, description, function, source_language)
+        mock_register_payload = json.load(f)
+    name = mock_register_payload["name"]
+    description = mock_register_payload["description"]
+    function = mock_register_payload["function"]
+    source_language = mock_register_payload["sourceLanguage"]
+    input_modality = mock_register_payload["input_modality"]
+    output_modality = mock_register_payload["output_modality"]
+    documentation_url = mock_register_payload["documentation_url"]
+    register_response = ModelFactory.create_asset_repo(name, description, function, source_language, input_modality, output_modality, documentation_url, config.TEAM_API_KEY)
     assert "id" in register_response.keys()
     assert "repositoryName" in register_response.keys()
     model_id = register_response["id"]
@@ -56,10 +56,7 @@ def test_create_and_upload_model():
     registry = login_response["registry"]
 
     # Push an image to ECR
-    # os.system("aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 535945872701.dkr.ecr.us-east-1.amazonaws.com")
     low_level_client = docker.APIClient(base_url="unix://var/run/docker.sock")
-    # low_level_client.pull("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash")
-    # low_level_client.tag("535945872701.dkr.ecr.us-east-1.amazonaws.com/bash", f"{registry}/{repo_name}")
     low_level_client.pull("bash")
     low_level_client.tag("bash", f"{registry}/{repo_name}")
     low_level_client.push(f"{registry}/{repo_name}", auth_config={"username": username, "password": password})

tests/image_upload_functional_test.py → tests/functional/model/image_upload_functional_test.py

@@ -1,13 +1,12 @@
 __author__ = "michaellam"
 from pathlib import Path
 import json
-from aixplain.utils import config
-from tests.test_utils import delete_asset, delete_service_account
 from aixplain.factories.model_factory import ModelFactory
+from tests.test_utils import delete_asset, delete_service_account
+from aixplain.utils import config
+import docker
 import pytest
 
-
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_login():
     response = ModelFactory.asset_repo_login()
     assert response["username"] == "AWS"
@@ -18,18 +17,17 @@ def test_login():
     delete_service_account(config.TEAM_API_KEY)
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_create_asset_repo():
     with open(Path("tests/test_requests/create_asset_request.json")) as f:
         mock_register_payload = json.load(f)
     name = mock_register_payload["name"]
-    host_machine = mock_register_payload["hostingMachine"]
-    version = mock_register_payload["version"]
     description = mock_register_payload["description"]
     function = mock_register_payload["function"]
     source_language = mock_register_payload["sourceLanguage"]
-    response = ModelFactory.create_asset_repo(name, host_machine, version, description, function, source_language)
-    print(response)
+    input_modality = mock_register_payload["input_modality"]
+    output_modality = mock_register_payload["output_modality"]
+    documentation_url = mock_register_payload["documentation_url"]
+    response = ModelFactory.create_asset_repo(name, description, function, source_language, input_modality, output_modality, documentation_url, config.TEAM_API_KEY)
     response_dict = dict(response)
     assert "id" in response_dict.keys()
     assert "repositoryName" in response_dict.keys()
@@ -38,7 +36,6 @@ def test_create_asset_repo():
     delete_asset(response["id"], config.TEAM_API_KEY)
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_list_host_machines():
     response = ModelFactory.list_host_machines()
     for hosting_machine_dict in response:
@@ -49,7 +46,6 @@ def test_list_host_machines():
         assert "hourlyCost" in hosting_machine_dict.keys()
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_get_functions():
     # Verbose
     response = ModelFactory.list_functions(True)

tests/test_requests/create_asset_request.json

@@ -1,8 +1,9 @@
 {
     "name": "mock_name",
-    "hostingMachine": "aix-2c-8g-od",
-    "version": "mock_version",
     "description": "mock_description",
-    "function": "Speech Recognition",
-    "sourceLanguage": "en"
+    "function": "Text Generation",
+    "sourceLanguage": "en",
+    "input_modality": "text",
+    "output_modality": "text",
+    "documentation_url": ""
 }

tests/image_upload_test.py → tests/unit/image_upload_test.py

@@ -13,7 +13,6 @@
 API_FIXED_HEADER = {"x-api-key": f"{config.TEAM_API_KEY}", "Content-Type": "application/json"}
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_login():
     url = urljoin(config.BACKEND_URL, f"sdk/ecr/#")
     with requests_mock.Mocker() as mock:
@@ -24,24 +23,26 @@ def test_login():
     assert creds == mock_json
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_create_asset_repo():
-    url_register = urljoin(config.BACKEND_URL, f"sdk/models/register")
+    url_register = urljoin(config.BACKEND_URL, f"sdk/models/onboard")
     url_function = urljoin(config.BACKEND_URL, f"sdk/functions")
+    print(f"URL_Register {url_register}")
     with requests_mock.Mocker() as mock:
         with open(Path("tests/mock_responses/create_asset_repo_response.json")) as f:
             mock_json_register = json.load(f)
-        mock.post(url_register, headers=API_FIXED_HEADER, json=mock_json_register)
+        mock.post(url_register, headers=API_FIXED_HEADER, json=mock_json_register, status_code=201)
+
         with open(Path("tests/mock_responses/list_functions_response.json")) as f:
             mock_json_functions = json.load(f)
         mock.get(url_function, headers=AUTH_FIXED_HEADER, json=mock_json_functions)
+
         model_id = ModelFactory.create_asset_repo(
-            "mock_name", "mock_machines", "mock_version", "mock_description", "Speech Recognition", "en", config.TEAM_API_KEY
+            "mock_name", "mock_description", "Text Generation", "en", "text", "text", api_key=config.TEAM_API_KEY
         )
+        # print(f"Model ID {model_id}")
     assert model_id == mock_json_register
 
 
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_list_host_machines():
     url = urljoin(config.BACKEND_URL, f"sdk/hosting-machines")
     with requests_mock.Mocker() as mock:
@@ -55,8 +56,6 @@ def test_list_host_machines():
         for key in machine_dict.keys():
             assert machine_dict[key] == mock_json_dict[key]
 
-
-@pytest.mark.skip(reason="Model Upload is deactivated for improvements.")
 def test_get_functions():
     url = urljoin(config.BACKEND_URL, f"sdk/functions")
     with requests_mock.Mocker() as mock:
@@ -66,7 +65,6 @@ def test_get_functions():
         functions = ModelFactory.list_functions(config.TEAM_API_KEY)
     assert functions == mock_json
 
-
 @pytest.mark.skip(reason="Not currently supported.")
 def test_list_image_repo_tags():
     model_id = "mock_id"