Parse port from base url (#2133)

* Parse port from base url

Signed-off-by: osamaahmed17 <[email protected]>

* Parse port from base url

Signed-off-by: osamaahmed17 <[email protected]>

* linting errors solved

Signed-off-by: osamaahmed17 <[email protected]>

* Update clients/python/src/model_registry/core.py

Co-authored-by: Jon Burdo <[email protected]>
Signed-off-by: Osama Tahir <[email protected]>

* added test case

Signed-off-by: osamaahmed17 <[email protected]>

* solve linting error

Signed-off-by: osamaahmed17 <[email protected]>

---------

Signed-off-by: osamaahmed17 <[email protected]>
Signed-off-by: Osama Tahir <[email protected]>
Co-authored-by: Jon Burdo <[email protected]>

Author: osamaahmed17

clients/python/src/model_registry/core.py

@@ -6,6 +6,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import TypeVar, cast, overload
+from urllib.parse import urlparse
 
 from mr_openapi import (
     ApiClient,
@@ -56,9 +57,12 @@ def secure_connection(
             user_token: The PEM-encoded user token as a string.
             custom_ca: The path to a PEM-
         """
+        parsed_url = urlparse(server_address)
+        host = server_address if parsed_url.port else f"{server_address}:{port}"
+
         return cls(
             Configuration(
-                f"{server_address}:{port}",
+                host=host,
                 access_token=user_token,
                 ssl_ca_cert=custom_ca,
             )
@@ -78,8 +82,11 @@ def insecure_connection(
             port: Server port.
             user_token: The PEM-encoded user token as a string.
         """
+        parsed_url = urlparse(server_address)
+        host = server_address if parsed_url.port else f"{server_address}:{port}"
+
         config = Configuration(
-            host=f"{server_address}:{port}",
+            host=host,
             access_token=user_token,
             verify_ssl=False,
         )

clients/python/tests/test_client.py

@@ -1178,8 +1178,7 @@ def test_user_token_from_envvar(monkeypatch, mock_get_registered_models):
 
     monkeypatch.setenv("KF_PIPELINES_SA_TOKEN_PATH", token_file_path)
     client = ModelRegistry(
-        server_address="http://localhost",
-        port=8080,
+        server_address="http://localhost:8080",
         author="test_author",
         is_secure=False,
         # user_token=None -> ... Let it read from Env var
@@ -1201,8 +1200,7 @@ def test_user_token_from_k8s_file(monkeypatch, mock_get_registered_models):
         k8s_token_file_path = k8s_token_file.name
     monkeypatch.setattr("model_registry._client.DEFAULT_K8S_SA_TOKEN_PATH", k8s_token_file_path)
     client = ModelRegistry(
-        server_address="http://localhost",
-        port=8080,
+        server_address="http://localhost:8080",
         author="test_author",
         is_secure=False,
         # user_token=None -> ... Let it read from K8s file
@@ -1228,8 +1226,7 @@ def test_user_token_envvar_priority_over_k8s(monkeypatch, mock_get_registered_mo
         k8s_token_file_path = k8s_token_file.name
     monkeypatch.setattr("model_registry._client.DEFAULT_K8S_SA_TOKEN_PATH", k8s_token_file_path)
     client = ModelRegistry(
-        server_address="http://localhost",
-        port=8080,
+        server_address="http://localhost:8080",
         author="test_author",
         is_secure=False,
         # user_token=None -> ... Let it read from Env var (and not from K8s file, given that Env var is set)
@@ -1248,8 +1245,7 @@ def test_user_token_missing_warning(monkeypatch, mock_get_registered_models):
     mock_warn = MagicMock()
     monkeypatch.setattr("model_registry._client.warn", mock_warn)
     client = ModelRegistry(
-        server_address="http://localhost",
-        port=8080,
+        server_address="http://localhost:8080",
         author="test_author",
         is_secure=False,
         # user_token=None -> ... Let it try to read, but missing Env var and K8s file, it will emit a warning

clients/python/tests/test_connection.py

@@ -0,0 +1,28 @@
+from model_registry.core import ModelRegistryAPIClient
+
+
+def test_connection_args():
+    """Test connection arguments parsing."""
+    # old-style: explicit port
+    server_address = "http://localhost"
+    port = 8080
+    host = f"{server_address}:{port}"
+
+    client = ModelRegistryAPIClient.insecure_connection(server_address, port)
+    assert client.config.host == host
+
+    # new-style: port in url
+    server_address_with_port = "http://localhost:9090"
+    client = ModelRegistryAPIClient.insecure_connection(server_address_with_port, port)
+    assert client.config.host == server_address_with_port
+
+    # Secure connection tests
+    client = ModelRegistryAPIClient.secure_connection(
+        server_address, port, user_token="token"  # noqa: S106
+    )
+    assert client.config.host == host
+
+    client = ModelRegistryAPIClient.secure_connection(
+        server_address_with_port, port, user_token="token"  # noqa: S106
+    )
+    assert client.config.host == server_address_with_port