Fix test_cli.py for snapshot create

Signed-off-by: Andre Kurait <[email protected]>

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/cli.py

@@ -94,6 +94,7 @@ def start_replayer_cmd(ctx):
 
 # ##################### SNAPSHOT ###################
 
+
 @cli.group(name="snapshot")
 @click.pass_obj
 def snapshot_group(ctx):
@@ -118,14 +119,16 @@ def create_snapshot_cmd(ctx):
 def status_snapshot_cmd(ctx):
     """Check the status of the snapshot"""
     snapshot = ctx.env.snapshot
-    _, message = logic_snapshot.status(snapshot, source_cluster=ctx.env.source_cluster, target_cluster=ctx.env.target_cluster)
+    _, message = logic_snapshot.status(snapshot, source_cluster=ctx.env.source_cluster,
+                                       target_cluster=ctx.env.target_cluster)
     click.echo(message)
 
 # ##################### BACKFILL ###################
 
 # As we add other forms of backfill migrations, we should incorporate a way to dynamically allow different sets of
 # arguments depending on the type of backfill migration
 
+
 @cli.group(name="backfill")
 @click.pass_obj
 def backfill_group(ctx):

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/environment.py

@@ -67,8 +67,8 @@ def __init__(self, config_file: str):
 
         if 'snapshot' in self.config:
             self.snapshot: Snapshot = S3Snapshot(self.config["snapshot"],
-                                                  source_cluster=self.source_cluster,
-                                                  target_cluster=self.target_cluster)
+                                                 source_cluster=self.source_cluster,
+                                                 target_cluster=self.target_cluster)
             logger.info(f"Snapshot initialized: {self.snapshot}")
         else:
             logger.info("No snapshot provided")

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/logic/snapshot.py

@@ -4,6 +4,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def create(snapshot: Snapshot, *args, **kwargs) -> Tuple[SnapshotStatus, str]:
     logger.info(f"Creating snapshot with {args=} and {kwargs=}")
     try:
@@ -16,6 +17,7 @@ def create(snapshot: Snapshot, *args, **kwargs) -> Tuple[SnapshotStatus, str]:
         return SnapshotStatus.COMPLETED, "Snapshot created successfully." + "\n" + result.value
     return SnapshotStatus.FAILED, "Snapshot creation failed." + "\n" + result.value
 
+
 def status(snapshot: Snapshot, *args, **kwargs) -> Tuple[SnapshotStatus, str]:
     logger.info("Getting snapshot status")
     try:

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/models/snapshot.py

@@ -2,7 +2,7 @@
 from enum import Enum
 import logging
 import subprocess
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional
 from console_link.models.cluster import Cluster
 from console_link.models.command_result import CommandResult
 from cerberus import Validator
@@ -14,8 +14,8 @@
         "NOT_STARTED",
         "RUNNING",
         "COMPLETED",
-        "FAILED"
-        ])
+        "FAILED"])
+
 
 class Snapshot(ABC):
     """
@@ -52,6 +52,7 @@ def status(self, *args, **kwargs) -> CommandResult:
     }
 }
 
+
 class S3Snapshot(Snapshot):
     def __init__(self, config: Dict, source_cluster: Cluster, target_cluster: Optional[Cluster] = None) -> None:
         super().__init__(config, source_cluster, target_cluster)
@@ -88,4 +89,4 @@ def create(self, *args, **kwargs) -> CommandResult:
             return CommandResult(success=False, value=f"Failed to create snapshot: {str(e)}")
 
     def status(self, *args, **kwargs) -> CommandResult:
-        return CommandResult(success=False, value=f"Command not implemented")
+        return CommandResult(success=False, value="Command not implemented")

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/console_link/models/utils.py

@@ -22,6 +22,7 @@ def raise_for_aws_api_error(response: Dict) -> None:
             status_code=status_code
         )
 
+
 class ExitCode(Enum):
     SUCCESS = 0
-    FAILURE = 1
+    FAILURE = 1

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/tests/data/services.yaml

@@ -16,3 +16,7 @@ metrics_source:
 backfill:
     reindex_from_snapshot:
         docker:
+snapshot:
+  snapshot_name: "test_snapshot"
+  s3_repo_uri: "s3://test-bucket"
+  s3_region: "us-east-2"

TrafficCapture/dockerSolution/src/main/docker/migrationConsole/lib/console_link/tests/test_cli.py

@@ -1,5 +1,5 @@
 import pathlib
-from console_link.models.command_result import CommandResult
+from console_link.models.snapshot import SnapshotStatus
 
 import requests_mock
 
@@ -68,32 +68,37 @@ def test_cli_with_backfill_describe(runner, env, mocker):
 
 
 def test_cli_snapshot_create(runner, env, mocker):
-    mock_create = mocker.patch('console_link.models.snapshot.Snapshot.create')
+    mock = mocker.patch('console_link.logic.snapshot.create')
 
-    # Set the mock return values
-    mock_create.return_value = CommandResult(success=True, value="Snapshot created successfully")
+    # Set the mock return value
+    mock.return_value = SnapshotStatus.COMPLETED, "Snapshot created successfully."
 
     # Test snapshot creation
-    result_create = runner.invoke(cli, ['--config-file', str(VALID_SERVICES_YAML), 'snapshot', 'create'], catch_exceptions=True)
-    assert result_create.exit_code == 0
-    assert "Snapshot created successfully" in result_create.output
+    result = runner.invoke(cli, ['--config-file', str(VALID_SERVICES_YAML), 'snapshot', 'create'],
+                           catch_exceptions=True)
+
+    assert result.exit_code == 0
+    assert "Snapshot created successfully" in result.output
 
     # Ensure the mocks were called
-    mock_create.assert_called_once()
+    mock.assert_called_once()
+
 
 @pytest.mark.skip(reason="Not implemented yet")
 def test_cli_snapshot_status(runner, env, mocker):
-    mock_status = mocker.patch('console_link.models.snapshot.Snapshot.status')
+    mock = mocker.patch('console_link.logic.snapshot.status')
+
+    # Set the mock return value
+    mock.return_value = SnapshotStatus.COMPLETED, "Snapshot status: COMPLETED"
 
-    mock_status.return_value = CommandResult(success=True, value="Snapshot status: COMPLETED")
-
     # Test snapshot status
-    result_status = runner.invoke(cli, ['--config-file', str(VALID_SERVICES_YAML), 'snapshot', 'status'], catch_exceptions=True)
-    assert result_status.exit_code == 0
-    assert "Snapshot status: COMPLETED" in result_status.output
+    result = runner.invoke(cli, ['--config-file', str(VALID_SERVICES_YAML), 'snapshot', 'status'],
+                           catch_exceptions=True)
+    assert result.exit_code == 0
+    assert "Snapshot status: COMPLETED" in result.output
 
     # Ensure the mocks were called
-    mock_status.assert_called_once()
+    mock.assert_called_once()
 
 
 source_cat_indices = """
@@ -125,4 +130,3 @@ def test_cli_cat_indices_e2e(runner, env):
     assert 'TARGET CLUSTER' in result.output
     assert source_cat_indices in result.output
     assert target_cat_indices in result.output
-