|  | 
| 24 | 24 | 
 | 
| 25 | 25 | package org.jenkinsci.plugins.workflow; | 
| 26 | 26 | 
 | 
|  | 27 | +import com.google.common.collect.ImmutableSet; | 
| 27 | 28 | import hudson.EnvVars; | 
| 28 | 29 | import hudson.model.EnvironmentContributor; | 
| 29 | 30 | import hudson.model.Run; | 
| 30 | 31 | import hudson.model.TaskListener; | 
| 31 | 32 | import java.io.IOException; | 
| 32 | 33 | import java.util.Collections; | 
|  | 34 | +import java.util.HashMap; | 
|  | 35 | +import java.util.Map; | 
| 33 | 36 | import java.util.Set; | 
|  | 37 | +import javax.annotation.CheckForNull; | 
|  | 38 | +import javax.annotation.Nonnull; | 
| 34 | 39 | import org.jenkinsci.plugins.workflow.cps.CpsFlowDefinition; | 
|  | 40 | +import org.jenkinsci.plugins.workflow.graph.FlowNode; | 
| 35 | 41 | import org.jenkinsci.plugins.workflow.job.WorkflowJob; | 
| 36 | 42 | import org.jenkinsci.plugins.workflow.job.WorkflowRun; | 
| 37 | 43 | import org.jenkinsci.plugins.workflow.steps.BodyExecutionCallback; | 
| 38 | 44 | import org.jenkinsci.plugins.workflow.steps.EnvironmentExpander; | 
| 39 | 45 | import org.jenkinsci.plugins.workflow.steps.Step; | 
| 40 | 46 | import org.jenkinsci.plugins.workflow.steps.StepContext; | 
| 41 | 47 | import org.jenkinsci.plugins.workflow.steps.StepDescriptor; | 
|  | 48 | +import org.jenkinsci.plugins.workflow.steps.StepEnvironmentContributor; | 
| 42 | 49 | import org.jenkinsci.plugins.workflow.steps.StepExecution; | 
|  | 50 | +import org.jenkinsci.plugins.workflow.steps.SynchronousStepExecution; | 
| 43 | 51 | import org.jenkinsci.plugins.workflow.test.steps.SemaphoreStep; | 
| 44 | 52 | import org.junit.ClassRule; | 
| 45 | 53 | import org.junit.Rule; | 
| @@ -145,4 +153,59 @@ private static class ExpanderImpl extends EnvironmentExpander { | 
| 145 | 153 |         } | 
| 146 | 154 |     } | 
| 147 | 155 | 
 | 
|  | 156 | +    @Issue("JENKINS-51170") | 
|  | 157 | +    @Test public void perStepEnvironment() { | 
|  | 158 | +        story.then(r -> { | 
|  | 159 | +            WorkflowJob p = r.jenkins.createProject(WorkflowJob.class, "p"); | 
|  | 160 | +            p.setDefinition(new CpsFlowDefinition("printEnv \"VAR\"; printEnv \"VAR\"", true)); | 
|  | 161 | +            WorkflowRun b = story.j.buildAndAssertSuccess(p); | 
|  | 162 | +            r.assertLogContains("VAR=1", b); | 
|  | 163 | +            r.assertLogContains("VAR=2", b); | 
|  | 164 | +        }); | 
|  | 165 | +    } | 
|  | 166 | + | 
|  | 167 | +    @TestExtension("perStepEnvironment") public static class StepEnvAdder extends StepEnvironmentContributor { | 
|  | 168 | +        private Map<String, Integer> stepNumbers = new HashMap<>(); | 
|  | 169 | + | 
|  | 170 | +        @Override | 
|  | 171 | +        public void buildEnvironmentFor(@CheckForNull StepContext stepContext, | 
|  | 172 | +            @Nonnull EnvVars envs, | 
|  | 173 | +            @CheckForNull TaskListener listener) throws IOException, InterruptedException { | 
|  | 174 | + | 
|  | 175 | +            FlowNode node = stepContext.get(FlowNode.class); | 
|  | 176 | +            int stepNumber = stepNumbers.computeIfAbsent(node.getId(), (k) -> stepNumbers.size() + 1); | 
|  | 177 | +            envs.override("VAR", String.valueOf(stepNumber)); | 
|  | 178 | +        } | 
|  | 179 | +    } | 
|  | 180 | + | 
|  | 181 | +    public static class PrintEnvStep extends Step { | 
|  | 182 | +        private final String var; | 
|  | 183 | +        @DataBoundConstructor | 
|  | 184 | +        public PrintEnvStep(String var) { | 
|  | 185 | +            this.var = var; | 
|  | 186 | +        } | 
|  | 187 | + | 
|  | 188 | +        @Override | 
|  | 189 | +        public StepExecution start(StepContext context) throws Exception { return new Execution(context, var); } | 
|  | 190 | + | 
|  | 191 | +        private static class Execution extends SynchronousStepExecution<Void> { | 
|  | 192 | +            private final String var; | 
|  | 193 | + | 
|  | 194 | +            Execution(StepContext context, String var) { super(context); this.var = var; } | 
|  | 195 | + | 
|  | 196 | +            @Override | 
|  | 197 | +            protected Void run() throws Exception { | 
|  | 198 | +                StepContext context = getContext(); | 
|  | 199 | +                String message = this.var + "=" + context.get(EnvVars.class).get(var); | 
|  | 200 | +                context.get(TaskListener.class).getLogger().println(message); | 
|  | 201 | +                return null; | 
|  | 202 | +            } | 
|  | 203 | +        } | 
|  | 204 | + | 
|  | 205 | +        @TestExtension("perStepEnvironment") public static class DescriptorImpl extends StepDescriptor { | 
|  | 206 | +            @Override public String getFunctionName() { return "printEnv"; } | 
|  | 207 | +            @Override public Set<? extends Class<?>> getRequiredContext() { return ImmutableSet.of(EnvVars.class, TaskListener.class); } | 
|  | 208 | +        } | 
|  | 209 | +    } | 
|  | 210 | + | 
| 148 | 211 | } | 
0 commit comments