Refactor

canonical · Mar 21, 2024 · a80f809 · a80f809
1 parent 6a93fb7
commit a80f809
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 15 deletions.
diff --git a/jupyter-pytorch-full/tests/test_access.py b/jupyter-pytorch-full/tests/test_access.py
@@ -3,13 +3,11 @@
 # See LICENSE file for licensing details.
 #
 #
-from pathlib import Path
-
 import subprocess
 import requests
 import tenacity
-import yaml
 
+from charmed_kubeflow_chisme.rock import CheckRock
 
 @tenacity.retry(
 stop=tenacity.stop_after_attempt(5),
@@ -22,9 +20,9 @@ def check_notebook_server_up(url):
 
 def main():
     """Test running container and imports."""
-    rock = yaml.safe_load(Path("rockcraft.yaml").read_text())
-    rock_image = rock["name"]
-    rock_version = rock["version"]
+    check_rock = CheckRock("rockcraft.yaml")
+    rock_image = check_rock.get_name()
+    rock_version = check_rock.get_version()
     LOCAL_ROCK_IMAGE = f"{rock_image}:{rock_version}"
 
     print(f"Running {LOCAL_ROCK_IMAGE}")

diff --git a/jupyter-pytorch-full/tests/test_imports.py b/jupyter-pytorch-full/tests/test_imports.py
@@ -3,18 +3,17 @@
 # See LICENSE file for licensing details.
 #
 #
-from pathlib import Path
-
 import os
 import subprocess
-import yaml
+
+from charmed_kubeflow_chisme.rock import CheckRock
 
 
 def main():
     """Test running container and imports."""
-    rock = yaml.safe_load(Path("rockcraft.yaml").read_text())
-    rock_image = rock["name"]
-    rock_version = rock["version"]
+    check_rock = CheckRock("rockcraft.yaml")
+    rock_image = check_rock.get_name()
+    rock_version = check_rock.get_version()
     LOCAL_ROCK_IMAGE = f"{rock_image}:{rock_version}"
 
     print(f"Running command in {LOCAL_ROCK_IMAGE}")

diff --git a/jupyter-pytorch-full/tests/test_rock.py b/jupyter-pytorch-full/tests/test_rock.py
@@ -2,11 +2,8 @@
 # See LICENSE file for licensing details.
 #
 #
-from pathlib import Path
-
 import pytest
 import subprocess
-import yaml
 
 from charmed_kubeflow_chisme.rock import CheckRock