Fix test_online_store_rest_client

logicalclocks · Jul 17, 2024 · 1a300fe · 1a300fe
1 parent 93b838f
commit 1a300fe
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/python/hopsworks_common/client/__init__.py b/python/hopsworks_common/client/__init__.py
@@ -18,7 +18,7 @@
 
 from typing import Literal, Optional, Union
 
-from hsfs.client import external, hopsworks
+from hopsworks_common.client import external, hopsworks
 
 
 _client: Union[hopsworks.Client, external.Client, None] = None

diff --git a/python/tests/core/test_online_store_rest_client.py b/python/tests/core/test_online_store_rest_client.py
@@ -13,10 +13,10 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 #
-import hsfs
+import hopsworks_common
 import pytest
 from furl import furl
-from hsfs.client import auth, exceptions, online_store_rest_client
+from hopsworks_common.client import auth, exceptions, online_store_rest_client
 
 
 class MockExternalClient:
@@ -50,13 +50,15 @@ def test_setup_rest_client_external(self, mocker, monkeypatch):
         def client_get_instance():
             return MockExternalClient()
 
-        monkeypatch.setattr(hsfs.client, "get_instance", client_get_instance)
+        monkeypatch.setattr(
+            hopsworks_common.client, "get_instance", client_get_instance
+        )
         variable_api_mock = mocker.patch(
-            "hsfs.core.variable_api.VariableApi.get_loadbalancer_external_domain",
+            "hopsworks_common.core.variable_api.VariableApi.get_loadbalancer_external_domain",
             return_value="app.hopsworks.ai",
         )
         ping_rdrs_mock = mocker.patch(
-            "hsfs.client.online_store_rest_client.OnlineStoreRestClientSingleton.is_connected",
+            "hopsworks_common.client.online_store_rest_client.OnlineStoreRestClientSingleton.is_connected",
         )
 
         # Act
@@ -86,14 +88,16 @@ def test_setup_online_store_rest_client_internal(self, mocker, monkeypatch):
         def client_get_instance():
             return MockInternalClient()
 
-        monkeypatch.setattr(hsfs.client, "get_instance", client_get_instance)
+        monkeypatch.setattr(
+            hopsworks_common.client, "get_instance", client_get_instance
+        )
         variable_api_mock = mocker.patch(
-            "hsfs.core.variable_api.VariableApi.get_service_discovery_domain",
+            "hopsworks_common.core.variable_api.VariableApi.get_service_discovery_domain",
             return_value="consul",
         )
         optional_config = {"api_key": "provided_api_key"}
         ping_rdrs_mock = mocker.patch(
-            "hsfs.client.online_store_rest_client.OnlineStoreRestClientSingleton.is_connected",
+            "hopsworks_common.client.online_store_rest_client.OnlineStoreRestClientSingleton.is_connected",
         )
 
         # Act