resolve comments

Azure · Jul 10, 2024 · 7405f58 · 7405f58
1 parent c94d15b
commit 7405f58
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 19 deletions.
diff --git a/src/acat/azext_acat/custom.py b/src/acat/azext_acat/custom.py
@@ -68,7 +68,7 @@ def _build_configuration(cls, ctx, credential, **kwargs):
             credential_scopes=resource_to_scopes(
                 ctx.cli_ctx.cloud.endpoints.active_directory_resource_id
             ),
-            **kwargs
+            **kwargs,
         )
 
 
@@ -149,16 +149,13 @@ def _build_arguments_schema(cls, *args, **kwargs):
         return args_schema
 
     def pre_operations(self):
-        cmd = OnboardAcat(cli_ctx=self.cli_ctx)
 
-        poller = cmd._handler(
-            {
-                "cmd": cmd,
+        poller = OnboardAcat(cli_ctx=self.cli_ctx)(
+            command_args={
                 "subscription_ids": ParseSubsFromResources(self.ctx.args.resources),
             }
         )
-        LongRunningOperation(cmd.cli_ctx, f"Starting {cmd.name}")(poller)
-
+        LongRunningOperation(self.cli_ctx)(poller)
         args = self.ctx.args
         args.time_zone = args.time_zone_with_default
         args.trigger_time = args.trigger_time_with_default
@@ -190,8 +187,6 @@ class DownloadAcatReportWithDupAadToken(_AcatDownloadSnapshot.SnapshotDownload):
 
     def _execute_operations(self):
         self.pre_operations()
-        if self.ctx.args.snapshot_name is None:
-            yield f"No snapshot found for report {self.ctx.args.report_name}"
         yield self.DownloadAcatReportWithDupAadToken(ctx=self.ctx)()
         self.post_operations()
 
@@ -243,11 +238,14 @@ def _output(self, *args, **kwargs):
 
     def pre_operations(self):
         args = self.ctx.args
-        cmd = GetControlAssessment(cli_ctx=self.cli_ctx)
-        result = cmd._handler(
-            {"cmd": cmd, "report_name": args.report_name, "compliance_status": "all"}
+        result = GetControlAssessment(cli_ctx=self.cli_ctx)(
+            command_args={"report_name": args.report_name, "compliance_status": "all"}
         )
         snapshot_name = result["snapshotName"] if result is not None else None
+        if snapshot_name is None:
+            raise ValueError(
+                f"No snapshot found for report {self.ctx.args.report_name}"
+            )
         args.snapshot_name = snapshot_name
 
 
@@ -284,7 +282,9 @@ def _output(self, *args, **kwargs):
             self.ctx.vars.instance.value, client_flatten=True
         )
         if len(snapshots) == 0:
-            return "No snapshot found"
+            raise ValueError(
+                f"No snapshot found for report {self.ctx.args.report_name}"
+            )
         latestSnapshot = snapshots[0]
         if self.ctx.args.compliance_status == "all":
             return latestSnapshot

diff --git a/src/acat/azext_acat/tests/latest/test_acat_report.py b/src/acat/azext_acat/tests/latest/test_acat_report.py
@@ -110,10 +110,6 @@ def test_get_control_assessments_no_snapshot(self):
         report = self.cmd(
             'acat report create --report-name {report_name_no_snapshot} --resources {resources}  --time-zone "UTC"    --trigger-time "2021-09-01T00:00:00Z"')
         assert(report.get_output_in_json().get('lastTriggerTime',None) == None)
-        ret=self.cmd(
-            'acat report get-control-assessments --report-name {report_name_no_snapshot}')
-        print(ret)
-        output=ret.get_output_in_json()
-        assert(output != None)
-        assert(output=='No snapshot found')
+        self.cmd(
+            'acat report get-control-assessments --report-name {report_name_no_snapshot}', expect_failure=True)
         self.cmd('acat report delete --report-name {report_name_no_snapshot} --yes')