[MISC] Data-replication Jubilant test fixes (#697)

Author: sinclert-canonical

tests/integration/high_availability/high_availability_helpers_new.py

@@ -3,10 +3,8 @@
 # See LICENSE file for licensing details.
 
 import json
-import os
 import subprocess
 from collections.abc import Callable
-from pathlib import Path
 
 import jubilant_backports
 from jubilant_backports import Juju
@@ -28,18 +26,6 @@
 JujuAppsStatusFn = Callable[[Status, str], bool]
 
 
-def _get_juju_keys() -> tuple[str, str]:
-    """Get Juju public and private keys."""
-    config_dir = Path("~/.local/share/juju")
-    if juju_data := os.getenv("JUJU_DATA"):
-        config_dir = Path(juju_data)
-
-    return (
-        str(config_dir.expanduser().resolve() / "ssh" / "juju_id_rsa.pub"),
-        str(config_dir.expanduser().resolve() / "ssh" / "juju_id_rsa"),
-    )
-
-
 def check_mysql_instances_online(juju: Juju, app_name: str) -> bool:
     """Checks whether all MySQL cluster instances are online.
 
@@ -121,20 +107,16 @@ def get_app_units(juju: Juju, app_name: str) -> dict[str, UnitStatus]:
 
 def scale_app_units(juju: Juju, app_name: str, num_units: int) -> None:
     """Scale a given application to a number of units."""
-    app_units = get_app_units(juju, app_name)
-    app_units_diff = len(app_units) - num_units
+    app_units = list(get_app_units(juju, app_name))
+    app_units_diff = num_units - len(app_units)
 
-    scale_func = None
     if app_units_diff > 0:
-        scale_func = juju.remove_unit
+        juju.add_unit(app_name, num_units=app_units_diff)
     if app_units_diff < 0:
-        scale_func = juju.add_unit
+        juju.remove_unit(*app_units[app_units_diff:])
     if app_units_diff == 0:
         return
 
-    for _ in range(abs(app_units_diff)):
-        scale_func(app_name, num_units=1)
-
     juju.wait(
         ready=lambda status: len(status.apps[app_name].units) == num_units,
         timeout=20 * MINUTE_SECS,
@@ -206,54 +188,6 @@ def get_unit_status_log(juju: Juju, unit_name: str, log_lines: int = 0) -> list[
     return json.loads(output)
 
 
-# TODO:
-#  Rely on Jubilant scp command once they make it easier
-#  for package users to deal with temporal directories
-#  https://github.com/canonical/jubilant/issues/201
-def scp_unit_file_from(
-    juju: Juju, app_name: str, unit_name: str, source_path: str, target_path: str
-) -> None:
-    """Copy a file from a given unit."""
-    juju_keys = _get_juju_keys()
-    unit_address = get_unit_ip(juju, app_name, unit_name)
-    unit_username = "ubuntu"
-
-    subprocess.check_output([
-        "scp",
-        "-B",
-        "-o",
-        "StrictHostKeyChecking=no",
-        "-i",
-        juju_keys[1],
-        f"{unit_username}@{unit_address}:{source_path}",
-        f"{target_path}",
-    ])
-
-
-# TODO:
-#  Rely on Jubilant scp command once they make it easier
-#  for package users to deal with temporal directories
-#  https://github.com/canonical/jubilant/issues/201
-def scp_unit_file_into(
-    juju: Juju, app_name: str, unit_name: str, source_path: str, target_path: str
-) -> None:
-    """Copy a file into a given unit."""
-    juju_keys = _get_juju_keys()
-    unit_address = get_unit_ip(juju, app_name, unit_name)
-    unit_username = "ubuntu"
-
-    subprocess.check_output([
-        "scp",
-        "-B",
-        "-o",
-        "StrictHostKeyChecking=no",
-        "-i",
-        juju_keys[1],
-        f"{source_path}",
-        f"{unit_username}@{unit_address}:{target_path}",
-    ])
-
-
 def get_relation_data(juju: Juju, app_name: str, rel_name: str) -> list[dict]:
     """Returns a list that contains the relation-data.
 

tests/integration/high_availability/test_replication_logs_rotation.py

@@ -3,6 +3,7 @@
 
 import logging
 import tempfile
+from pathlib import Path
 
 import jubilant_backports
 import pytest
@@ -18,8 +19,6 @@
 from .high_availability_helpers_new import (
     get_app_leader,
     get_unit_process_id,
-    scp_unit_file_from,
-    scp_unit_file_into,
     wait_for_apps_status,
 )
 
@@ -91,7 +90,6 @@ def test_log_rotation(juju: Juju) -> None:
         logging.info("Writing some data to the text log files")
         write_unit_file(
             juju=juju,
-            app_name=MYSQL_APP_NAME,
             unit_name=mysql_app_leader,
             file_path=f"{mysql_logs_path}/{log_type}.log",
             file_data=f"{log_type} content",
@@ -118,7 +116,6 @@ def test_log_rotation(juju: Juju) -> None:
     for log_type in log_types:
         active_log_file_data = read_unit_file(
             juju=juju,
-            app_name=MYSQL_APP_NAME,
             unit_name=mysql_app_leader,
             file_path=f"{mysql_logs_path}/{log_type}.log",
         )
@@ -132,7 +129,6 @@ def test_log_rotation(juju: Juju) -> None:
         archive_log_file_name = archive_log_files_listed[0].split()[-1]
         archive_log_file_data = read_unit_file(
             juju=juju,
-            app_name=MYSQL_APP_NAME,
             unit_name=mysql_app_leader,
             file_path=f"{archive_log_dir}/{archive_log_file_name}",
         )
@@ -174,12 +170,11 @@ def list_unit_files(juju: Juju, unit_name: str, file_path: str) -> list[str]:
     ]
 
 
-def read_unit_file(juju: Juju, app_name: str, unit_name: str, file_path: str) -> str:
+def read_unit_file(juju: Juju, unit_name: str, file_path: str) -> str:
     """Read contents from file in the provided unit.
 
     Args:
         juju: The Juju instance
-        app_name: The name of the app
         unit_name: The name of the unit to read the file from
         file_path: The path of the unit to read the file
     """
@@ -188,42 +183,35 @@ def read_unit_file(juju: Juju, app_name: str, unit_name: str, file_path: str) ->
     juju.exec(f"sudo cp {file_path} {temp_path}", unit=unit_name)
     juju.exec(f"sudo chown ubuntu:ubuntu {temp_path}", unit=unit_name)
 
-    with tempfile.NamedTemporaryFile(mode="r+") as temp_file:
-        scp_unit_file_from(
-            juju=juju,
-            app_name=app_name,
-            unit_name=unit_name,
-            source_path=temp_path,
-            target_path=temp_file.name,
+    with tempfile.NamedTemporaryFile(mode="r+", dir=Path.home()) as temp_file:
+        juju.scp(
+            f"{unit_name}:{temp_path}",
+            f"{temp_file.name}",
         )
         contents = temp_file.read()
 
     juju.exec(f"sudo rm {temp_path}", unit=unit_name)
     return contents
 
 
-def write_unit_file(juju: Juju, app_name: str, unit_name: str, file_path: str, file_data: str):
+def write_unit_file(juju: Juju, unit_name: str, file_path: str, file_data: str):
     """Write content to the file in the provided unit.
 
     Args:
         juju: The Juju instance
-        app_name: The name of the app
         unit_name: The name of the unit to write the file into
         file_path: The path of the unit to write the file
         file_data: The data to write to the file.
     """
     temp_path = "/tmp/file"
 
-    with tempfile.NamedTemporaryFile(mode="w") as temp_file:
+    with tempfile.NamedTemporaryFile(mode="w", dir=Path.home()) as temp_file:
         temp_file.write(file_data)
         temp_file.flush()
 
-        scp_unit_file_into(
-            juju=juju,
-            app_name=app_name,
-            unit_name=unit_name,
-            source_path=temp_file.name,
-            target_path=temp_path,
+        juju.scp(
+            f"{temp_file.name}",
+            f"{unit_name}:{temp_path}",
         )
 
     juju.exec(f"sudo mv {temp_path} {file_path}", unit=unit_name)