Add support for exit codes

- New configurations to CommandRegistration
- Re-using exit code concepts from boot
- Handling exit codes only in non-interactive mode
- Adding e2e commands and tests for better coverage
- Fixes #431

spring-shell-core/src/main/java/org/springframework/shell/Shell.java

@@ -15,12 +15,12 @@
  */
 package org.springframework.shell;
 
-import java.io.IOException;
 import java.lang.reflect.UndeclaredThrowableException;
 import java.nio.channels.ClosedByInterruptException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import javax.validation.Validator;
@@ -42,6 +42,9 @@ import org.springframework.shell.command.CommandExecution.CommandExecutionExcept
 import org.springframework.shell.command.CommandExecution.CommandExecutionHandlerMethodArgumentResolvers;
 import org.springframework.shell.command.CommandRegistration;
 import org.springframework.shell.completion.CompletionResolver;
+import org.springframework.shell.context.InteractionMode;
+import org.springframework.shell.context.ShellContext;
+import org.springframework.shell.exit.ExitCodeMappings;
 
 /**
  * Main class implementing a shell loop.
@@ -65,6 +68,8 @@ public class Shell {
 	protected List<CompletionResolver> completionResolvers = new ArrayList<>();
 	private CommandExecutionHandlerMethodArgumentResolvers argumentResolvers;
 	private ConversionService conversionService = new DefaultConversionService();
+	private final ShellContext shellContext;
+	private final ExitCodeMappings exitCodeMappings;
 
 	/**
 	 * Marker object to distinguish unresolved arguments from {@code null}, which is a valid
@@ -74,10 +79,13 @@ public class Shell {
 
 	private Validator validator = Utils.defaultValidator();
 
-	public Shell(ResultHandlerService resultHandlerService, CommandCatalog commandRegistry, Terminal terminal) {
+	public Shell(ResultHandlerService resultHandlerService, CommandCatalog commandRegistry, Terminal terminal,
+			ShellContext shellContext, ExitCodeMappings exitCodeMappings) {
 		this.resultHandlerService = resultHandlerService;
 		this.commandRegistry = commandRegistry;
 		this.terminal = terminal;
+		this.shellContext = shellContext;
+		this.exitCodeMappings = exitCodeMappings;
 	}
 
 	@Autowired
@@ -108,7 +116,7 @@ public class Shell {
 	 * (<em>e.g.</em> a {@literal script} command).
 	 * </p>
 	 */
-	public void run(InputProvider inputProvider) throws IOException {
+	public void run(InputProvider inputProvider) throws Exception {
 		Object result = null;
 		while (!(result instanceof ExitRequest)) { // Handles ExitRequest thrown from Quit command
 			Input input;
@@ -130,6 +138,18 @@ public class Shell {
 			if (result != NO_INPUT && !(result instanceof ExitRequest)) {
 				resultHandlerService.handle(result);
 			}
+
+			// throw if not in interactive mode so that boot's exit code feature
+			// can contribute exit code. we can't throw when in interactive mode as
+			// that would exit a shell
+			if (this.shellContext != null && this.shellContext.getInteractionMode() != InteractionMode.INTERACTIVE) {
+				if (result instanceof CommandExecution.CommandParserExceptionsException) {
+					throw (CommandExecution.CommandParserExceptionsException) result;
+				}
+				else if (result instanceof Exception) {
+					throw (Exception) result;
+				}
+			}
 		}
 	}
 
@@ -169,6 +189,12 @@ public class Shell {
 				.findFirst();
 
 			if (commandRegistration.isPresent()) {
+				if (this.exitCodeMappings != null) {
+					List<Function<Throwable, Integer>> mappingFunctions = commandRegistration.get().getExitCode()
+							.getMappingFunctions();
+					this.exitCodeMappings.reset(mappingFunctions);
+				}
+
 				List<String> wordsForArgs = wordsForArguments(command, words);
 
 				Thread commandThread = Thread.currentThread();

spring-shell-core/src/main/java/org/springframework/shell/command/CommandExitCode.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.shell.command;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * Interface representing an exit code in a command.
+ *
+ * @author Janne Valkealahti
+ */
+public interface CommandExitCode {
+
+	/**
+	 * Gets a function mappings from exceptions to exit codes.
+	 *
+	 * @return function mappings
+	 */
+	List<Function<Throwable, Integer>> getMappingFunctions();
+
+	/**
+	 * Gets an instance of a default {@link CommandExitCode}.
+	 *
+	 * @return a command exit code
+	 */
+	public static CommandExitCode of() {
+		return of(new ArrayList<>());
+	}
+
+	/**
+	 * Gets an instance of a default {@link CommandExitCode}.
+	 *
+	 * @param functions the function mappings
+	 * @return a command exit code
+	 */
+	public static CommandExitCode of(List<Function<Throwable, Integer>> functions) {
+		return new DefaultCommandExitCode(functions);
+	}
+
+	static class DefaultCommandExitCode implements CommandExitCode {
+
+		private final List<Function<Throwable, Integer>> functions;
+
+		DefaultCommandExitCode(	List<Function<Throwable, Integer>> functions) {
+			this.functions = functions;
+		}
+
+		@Override
+		public List<Function<Throwable, Integer>> getMappingFunctions() {
+			return functions;
+		}
+	}
+}

spring-shell-core/src/main/java/org/springframework/shell/command/CommandRegistration.java

@@ -98,6 +98,13 @@ public interface CommandRegistration {
 	 */
 	List<CommandAlias> getAliases();
 
+	/**
+	 * Gets an exit code.
+	 *
+	 * @return the exit code
+	 */
+	CommandExitCode getExitCode();
+
 	/**
 	 * Gets a new instance of a {@link Buidler}.
 	 *
@@ -386,6 +393,35 @@ public interface CommandRegistration {
 		Builder and();
 	}
 
+	/**
+	 * Spec defining an exit code.
+	 */
+	public interface ExitCodeSpec {
+
+		/**
+		 * Define mapping from exception to code.
+		 *
+		 * @param e the exception
+		 * @param code the exit code
+		 * @return a target spec for chaining
+		 */
+		ExitCodeSpec map(Class<? extends Throwable> e, int code);
+
+		/**
+		 *
+		 * @param function
+		 * @return
+		 */
+		ExitCodeSpec map(Function<Throwable, Integer> function);
+
+		/**
+		 * Return a builder for chaining.
+		 *
+		 * @return a builder for chaining
+		 */
+		Builder and();
+	}
+
 	/**
 	 * Builder interface for {@link CommandRegistration}.
 	 */
@@ -456,6 +492,13 @@ public interface CommandRegistration {
 		 */
 		AliasSpec withAlias();
 
+		/**
+		 * Define an exit code what this command should execute
+		 *
+		 * @return exit code spec for chaining
+		 */
+		ExitCodeSpec withExitCode();
+
 		/**
 		 * Builds a {@link CommandRegistration}.
 		 *
@@ -689,6 +732,39 @@ public interface CommandRegistration {
 		}
 	}
 
+	static class DefaultExitCodeSpec implements ExitCodeSpec {
+
+		private BaseBuilder builder;
+		private final List<Function<Throwable, Integer>> functions = new ArrayList<>();
+
+		DefaultExitCodeSpec(BaseBuilder builder) {
+			this.builder = builder;
+		}
+
+		@Override
+		public ExitCodeSpec map(Class<? extends Throwable> e, int code) {
+			Function<Throwable, Integer> f = t -> {
+				if (ObjectUtils.nullSafeEquals(t.getClass(), e)) {
+					return code;
+				}
+				return 0;
+			};
+			this.functions.add(f);
+			return this;
+		}
+
+		@Override
+		public ExitCodeSpec map(Function<Throwable, Integer> function) {
+			this.functions.add(function);
+			return this;
+		}
+
+		@Override
+		public Builder and() {
+			return builder;
+		}
+	}
+
 	static class DefaultCommandRegistration implements CommandRegistration {
 
 		private String command;
@@ -699,10 +775,11 @@ public interface CommandRegistration {
 		private List<DefaultOptionSpec> optionSpecs;
 		private DefaultTargetSpec targetSpec;
 		private List<DefaultAliasSpec> aliasSpecs;
+		private DefaultExitCodeSpec exitCodeSpec;
 
 		public DefaultCommandRegistration(String[] commands, InteractionMode interactionMode, String group,
 				String description, Supplier<Availability> availability, List<DefaultOptionSpec> optionSpecs,
-				DefaultTargetSpec targetSpec, List<DefaultAliasSpec> aliasSpecs) {
+				DefaultTargetSpec targetSpec, List<DefaultAliasSpec> aliasSpecs, DefaultExitCodeSpec exitCodeSpec) {
 			this.command = commandArrayToName(commands);
 			this.interactionMode = interactionMode;
 			this.group = group;
@@ -711,6 +788,7 @@ public interface CommandRegistration {
 			this.optionSpecs = optionSpecs;
 			this.targetSpec = targetSpec;
 			this.aliasSpecs = aliasSpecs;
+			this.exitCodeSpec = exitCodeSpec;
 		}
 
 		@Override
@@ -769,6 +847,16 @@ public interface CommandRegistration {
 				.collect(Collectors.toList());
 		}
 
+		@Override
+		public CommandExitCode getExitCode() {
+			if (this.exitCodeSpec == null) {
+				return CommandExitCode.of();
+			}
+			else {
+				return CommandExitCode.of(exitCodeSpec.functions);
+			}
+		}
+
 		private static String commandArrayToName(String[] commands) {
 			return Arrays.asList(commands).stream()
 				.flatMap(c -> Stream.of(c.split(" ")))
@@ -792,6 +880,7 @@ public interface CommandRegistration {
 		private List<DefaultOptionSpec> optionSpecs = new ArrayList<>();
 		private List<DefaultAliasSpec> aliasSpecs = new ArrayList<>();
 		private DefaultTargetSpec targetSpec;
+		private DefaultExitCodeSpec exitCodeSpec;
 
 		@Override
 		public Builder command(String... commands) {
@@ -848,7 +937,14 @@ public interface CommandRegistration {
 			DefaultAliasSpec spec = new DefaultAliasSpec(this);
 			this.aliasSpecs.add(spec);
 			return spec;
-		};
+		}
+
+		@Override
+		public ExitCodeSpec withExitCode() {
+			DefaultExitCodeSpec spec = new DefaultExitCodeSpec(this);
+			this.exitCodeSpec = spec;
+			return spec;
+		}
 
 		@Override
 		public CommandRegistration build() {
@@ -856,7 +952,7 @@ public interface CommandRegistration {
 			Assert.notNull(targetSpec, "target cannot be empty");
 			Assert.state(!(targetSpec.bean != null && targetSpec.function != null), "only one target can exist");
 			return new DefaultCommandRegistration(commands, interactionMode, group, description, availability,
-					optionSpecs, targetSpec, aliasSpecs);
+					optionSpecs, targetSpec, aliasSpecs, exitCodeSpec);
 		}
 	}
 }

spring-shell-core/src/main/java/org/springframework/shell/exit/ExitCodeMappings.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.shell.exit;
+
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * Interface used with implementation of a boot's ExitCodeExceptionMapper
+ * in a context of spring-shell spesific one. Mostly needed not to have a
+ * direct dependencies to boot classes as currently only one implementation
+ * instance can exist which we need to reset between command executions.
+ *
+ * @author Janne Valkealahti
+ */
+public interface ExitCodeMappings {
+
+	/**
+	 * Reset mappings into a given functions.
+	 *
+	 * @param functions the mapping functions
+	 */
+	void reset(List<Function<Throwable, Integer>> functions);
+}

spring-shell-core/src/test/java/org/springframework/shell/ShellTests.java

@@ -15,7 +15,6 @@
  */
 package org.springframework.shell;
 
-import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
@@ -71,7 +70,7 @@ public class ShellTests {
 	}
 
 	@Test
-	public void commandMatch() throws IOException {
+	public void commandMatch() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "hello world how are you doing ?");
 		doThrow(new Exit()).when(resultHandlerService).handle(any());
 
@@ -97,7 +96,7 @@ public class ShellTests {
 	}
 
 	@Test
-	public void commandNotFound() throws IOException {
+	public void commandNotFound() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "hello world how are you doing ?");
 		doThrow(new Exit()).when(resultHandlerService).handle(isA(CommandNotFound.class));
 
@@ -122,7 +121,7 @@ public class ShellTests {
 
 	@Test
 	// See https://github.com/spring-projects/spring-shell/issues/142
-	public void commandNotFoundPrefix() throws IOException {
+	public void commandNotFoundPrefix() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "helloworld how are you doing ?");
 		doThrow(new Exit()).when(resultHandlerService).handle(isA(CommandNotFound.class));
 
@@ -146,7 +145,7 @@ public class ShellTests {
 	}
 
 	@Test
-	public void noCommand() throws IOException {
+	public void noCommand() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "", () -> "hello world how are you doing ?", null);
 		doThrow(new Exit()).when(resultHandlerService).handle(any());
 
@@ -172,7 +171,7 @@ public class ShellTests {
 	}
 
 	@Test
-	public void commandThrowingAnException() throws IOException {
+	public void commandThrowingAnException() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "fail");
 		doThrow(new Exit()).when(resultHandlerService).handle(isA(SomeException.class));
 
@@ -199,7 +198,7 @@ public class ShellTests {
 	}
 
 	@Test
-	public void comments() throws IOException {
+	public void comments() throws Exception {
 		when(inputProvider.readInput()).thenReturn(() -> "// This is a comment", (Input) null);
 
 		shell.run(inputProvider);

spring-shell-core/src/test/java/org/springframework/shell/command/CommandRegistrationTests.java

@@ -336,7 +336,7 @@ public class CommandRegistrationTests extends AbstractCommandTests {
 			assertThat(registration.getOptions()).hasSize(1);
 			assertThat(registration.getOptions().get(0).getArityMin()).isEqualTo(0);
 			assertThat(registration.getOptions().get(0).getArityMax()).isEqualTo(0);
-		}
+	}
 
 	@Test
 	public void testArityViaEnum() {
@@ -353,8 +353,7 @@ public class CommandRegistrationTests extends AbstractCommandTests {
 			assertThat(registration.getOptions()).hasSize(1);
 			assertThat(registration.getOptions().get(0).getArityMin()).isEqualTo(0);
 			assertThat(registration.getOptions().get(0).getArityMax()).isEqualTo(0);
-		}
-
+	}
 
 	@Test
 	public void testAliases() {
@@ -381,4 +380,32 @@ public class CommandRegistrationTests extends AbstractCommandTests {
 		assertThat(registration.getAliases().get(0).getGroup()).isEqualTo("Alias Group");
 		assertThat(registration.getAliases().get(1).getGroup()).isEqualTo("Alias Group");
 	}
+
+	@Test
+	public void testExitCodes() {
+		CommandRegistration registration;
+		registration = CommandRegistration.builder()
+			.command("command1")
+			.withTarget()
+				.function(function1)
+				.and()
+			.build();
+		assertThat(registration.getExitCode()).isNotNull();
+		assertThat(registration.getExitCode().getMappingFunctions()).hasSize(0);
+
+		registration = CommandRegistration.builder()
+			.command("command1")
+			.withExitCode()
+				.map(RuntimeException.class, 1)
+				.map(IllegalArgumentException.class, 2)
+				.map(e -> 1)
+				.map(e -> 2)
+				.and()
+			.withTarget()
+				.function(function1)
+				.and()
+			.build();
+		assertThat(registration.getExitCode()).isNotNull();
+		assertThat(registration.getExitCode().getMappingFunctions()).hasSize(4);
+	}
 }