Repository: apache/gravitino (Java)

Issue:
[Improvement] Improve validate in UpdateModelVersionComment.java and UpdateModelVersionAliases.java
### What would you like to be improved?

Add a check to stop potential RTE.

### How should we improve?

Add a check to make sure the alias and version are not null, just like what is done in UpdateModelVersionUri.java. Also add unit tests.


Discussion:
May I work on this issue? Please let me know if there are any special considerations or things I should keep in mind.
Nothing special, other than we use Spotless to format code, and you should add a unit test to confirm any changes.
Thank you. I’ll work on it and keep you updated.



Patch:
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java
index 9c36a14cf19..c6385252ec7 100644
--- a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java
+++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionComment.java
@@ -116,7 +116,9 @@ public Command validate() {
     if (alias != null && version != null) {
       exitWithError("Cannot specify both alias and version");
     }
-
+    if (alias == null && version == null) {
+      exitWithError("Either alias or version must be specified");
+    }
     return this;
   }
 }
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java
index 2e5a00e01e0..7b0ea582496 100644
--- a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java
+++ b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/UpdateModelVersionUri.java
@@ -116,7 +116,9 @@ public Command validate() {
     if (alias != null && version != null) {
       exitWithError("Cannot specify both alias and version");
     }
-
+    if (alias == null && version == null) {
+      exitWithError("Either alias or version must be specified");
+    }
     return this;
   }
 }


Test patch:
diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionAliases.java b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionAliases.java
new file mode 100644
index 00000000000..2feb5fe498e
--- /dev/null
+++ b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionAliases.java
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.cli.commands;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.gravitino.Catalog;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.cli.CommandContext;
+import org.apache.gravitino.cli.Main;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.model.ModelCatalog;
+import org.apache.gravitino.model.ModelVersionChange;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+class TestUpdateModelVersionAliases {
+
+  @BeforeEach
+  void setUp() {
+    Main.useExit = false;
+  }
+
+  @AfterEach
+  void tearDown() {
+    Main.useExit = true;
+  }
+
+  @Test
+  void handleWithAliasCallsAlterModelVersion() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        Mockito.spy(
+            new UpdateModelVersionAliases(
+                context,
+                "metalake1",
+                "catalog1",
+                "schema1",
+                "model1",
+                null,
+                "alias1",
+                new String[] {"add1", "add2"},
+                new String[] {"remove1"}));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            ArgumentMatchers.eq("alias1"),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command)
+        .printInformation(ArgumentMatchers.contains("alias alias1 aliases changed."));
+  }
+
+  @Test
+  void handleWithVersionCallsAlterModelVersion() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        Mockito.spy(
+            new UpdateModelVersionAliases(
+                context,
+                "metalake1",
+                "catalog1",
+                "schema1",
+                "model1",
+                99,
+                null,
+                new String[] {"add1"},
+                new String[] {"remove1"}));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            ArgumentMatchers.eq(99),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command).printInformation(Mockito.contains("version 99 aliases changed."));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNotNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        new UpdateModelVersionAliases(
+            context,
+            "metalake1",
+            "catalog1",
+            "schema1",
+            "model1",
+            1,
+            "alias1",
+            new String[] {"aliasA"},
+            new String[] {"aliasB"});
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        new UpdateModelVersionAliases(
+            context,
+            "metalake1",
+            "catalog1",
+            "schema1",
+            "model1",
+            null,
+            null,
+            new String[] {"aliasA"},
+            new String[] {"aliasB"});
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateOnlyAliasSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        new UpdateModelVersionAliases(
+            context,
+            "metalake1",
+            "catalog1",
+            "schema1",
+            "model1",
+            null,
+            "alias1",
+            new String[] {"aliasA"},
+            new String[] {"aliasB"});
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+
+  @Test
+  void validateOnlyVersionSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionAliases command =
+        new UpdateModelVersionAliases(
+            context,
+            "metalake1",
+            "catalog1",
+            "schema1",
+            "model1",
+            1,
+            null,
+            new String[] {"aliasA"},
+            new String[] {"aliasB"});
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+}
diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionComment.java b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionComment.java
new file mode 100644
index 00000000000..6ffd4d82fee
--- /dev/null
+++ b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionComment.java
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.cli.commands;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.gravitino.Catalog;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.cli.CommandContext;
+import org.apache.gravitino.cli.Main;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.model.ModelCatalog;
+import org.apache.gravitino.model.ModelVersionChange;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+class TestUpdateModelVersionComment {
+
+  @BeforeEach
+  void setUp() {
+    Main.useExit = false;
+  }
+
+  @AfterEach
+  void tearDown() {
+    Main.useExit = true;
+  }
+
+  @Test
+  void handleWithAliasCallsAlterModelVersionWithAlias() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        Mockito.spy(
+            new UpdateModelVersionComment(
+                context,
+                "metalake1",
+                "catalog1",
+                "schema1",
+                "model1",
+                null,
+                "alias1",
+                "new comment"));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            Mockito.eq("alias1"),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command).printInformation(Mockito.contains("alias alias1 comment changed."));
+  }
+
+  @Test
+  void handleWithVersionCallsAlterModelVersionWithVersion() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        Mockito.spy(
+            new UpdateModelVersionComment(
+                context, "metalake1", "catalog1", "schema1", "model1", 2, null, "new comment"));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            Mockito.eq(2),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command).printInformation(Mockito.contains("version 2 comment changed."));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNotNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        new UpdateModelVersionComment(
+            context, "metalake1", "catalog1", "schema1", "model1", 1, "alias1", "comment");
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        new UpdateModelVersionComment(
+            context, "metalake1", "catalog1", "schema1", "model1", null, null, "comment");
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateOnlyAliasSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        new UpdateModelVersionComment(
+            context, "metalake1", "catalog1", "schema1", "model1", null, "alias1", "comment");
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+
+  @Test
+  void validateOnlyVersionSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionComment command =
+        new UpdateModelVersionComment(
+            context, "metalake1", "catalog1", "schema1", "model1", 1, null, "comment");
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+}
diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionUri.java b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionUri.java
new file mode 100644
index 00000000000..c52bc518a1d
--- /dev/null
+++ b/clients/cli/src/test/java/org/apache/gravitino/cli/commands/TestUpdateModelVersionUri.java
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.cli.commands;
+
+import static org.mockito.Mockito.spy;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.gravitino.Catalog;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.cli.CommandContext;
+import org.apache.gravitino.cli.Main;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.model.ModelCatalog;
+import org.apache.gravitino.model.ModelVersionChange;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+class TestUpdateModelVersionUri {
+
+  @BeforeEach
+  void setUp() {
+    Main.useExit = false;
+  }
+
+  @AfterEach
+  void tearDown() {
+    Main.useExit = true;
+  }
+
+  @Test
+  void handleWithAliasCallsAlterModelVersion() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        spy(
+            new UpdateModelVersionUri(
+                context,
+                "metalake1",
+                "catalog1",
+                "schema1",
+                "model1",
+                null,
+                "alias1",
+                "http://some.uri"));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            ArgumentMatchers.eq("alias1"),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command)
+        .printInformation(ArgumentMatchers.contains("alias alias1 uri changed."));
+  }
+
+  @Test
+  void handleWithVersionCallsAlterModelVersion() throws Exception {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        spy(
+            new UpdateModelVersionUri(
+                context, "metalake1", "catalog1", "schema1", "model1", 4, null, "http://some.uri"));
+
+    GravitinoClient mockClient = Mockito.mock(GravitinoClient.class);
+    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    ModelCatalog mockModelCatalog = Mockito.mock(ModelCatalog.class);
+
+    Mockito.doReturn(mockClient).when(command).buildClient("metalake1");
+    Mockito.when(mockClient.loadCatalog("catalog1")).thenReturn(mockCatalog);
+    Mockito.when(mockCatalog.asModelCatalog()).thenReturn(mockModelCatalog);
+
+    command.handle();
+
+    Mockito.verify(mockModelCatalog)
+        .alterModelVersion(
+            ArgumentMatchers.any(NameIdentifier.class),
+            ArgumentMatchers.eq(4),
+            ArgumentMatchers.any(ModelVersionChange.class));
+    Mockito.verify(command).printInformation(ArgumentMatchers.contains("version 4 uri changed."));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNotNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        new UpdateModelVersionUri(
+            context, "metalake1", "catalog1", "schema1", "model1", 1, "alias1", "http://new.uri");
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateBothAliasAndVersionNullShouldExitWithError() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        new UpdateModelVersionUri(
+            context, "metalake1", "catalog1", "schema1", "model1", null, null, "http://new.uri");
+
+    RuntimeException ex = Assertions.assertThrows(RuntimeException.class, command::validate);
+    Assertions.assertTrue(ex.getMessage().contains("Exit with code"));
+  }
+
+  @Test
+  void validateOnlyAliasSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        new UpdateModelVersionUri(
+            context,
+            "metalake1",
+            "catalog1",
+            "schema1",
+            "model1",
+            null,
+            "alias1",
+            "http://new.uri");
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+
+  @Test
+  void validateOnlyVersionSetShouldPass() {
+    CommandLine mockCmdLine = Mockito.mock(CommandLine.class);
+    CommandContext context = new CommandContext(mockCmdLine);
+
+    UpdateModelVersionUri command =
+        new UpdateModelVersionUri(
+            context, "metalake1", "catalog1", "schema1", "model1", 1, null, "http://new.uri");
+
+    Assertions.assertDoesNotThrow(command::validate);
+  }
+}