mock-tasks: allow matching mocks by meta attributes or step name (#1024)

plugins/tasks/mock/src/main/java/com/walmartlabs/concord/plugins/mock/MockDefinition.java

@@ -35,6 +35,14 @@ public MockDefinition(Map<String, Object> definition) {
         this.definition = definition;
     }
 
+    public String stepName() {
+        return MapUtils.getString(definition, "stepName");
+    }
+
+    public Map<String, Object> stepMeta() {
+        return MapUtils.getMap(definition, "stepMeta", Map.of());
+    }
+
     public String task() {
         try {
             return MapUtils.assertString(definition, "task");
@@ -66,6 +74,4 @@ public Serializable result() {
     public String throwError() {
         return MapUtils.getString(definition, "throwError");
     }
-
-
 }

plugins/tasks/mock/src/main/java/com/walmartlabs/concord/plugins/mock/MockDefinitionContext.java

@@ -0,0 +1,43 @@
+package com.walmartlabs.concord.plugins.mock;
+
+/*-
+ * *****
+ * Concord
+ * -----
+ * Copyright (C) 2017 - 2024 Walmart Inc.
+ * -----
+ * 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
+ * 
+ *      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.
+ * =====
+ */
+
+import com.walmartlabs.concord.runtime.v2.model.Step;
+import com.walmartlabs.concord.runtime.v2.sdk.Variables;
+
+import java.util.Objects;
+
+public record MockDefinitionContext(Step currentStep, String taskName, Variables input, String method, Object[] params) {
+
+    public static MockDefinitionContext task(Step currentStep, String taskName, Variables input) {
+        Objects.requireNonNull(taskName);
+        Objects.requireNonNull(input);
+
+        return new MockDefinitionContext(currentStep, taskName, input, null, null);
+    }
+
+    public static MockDefinitionContext method(Step currentStep, String taskName, String methodName, Object[] params) {
+        Objects.requireNonNull(taskName);
+        Objects.requireNonNull(methodName);
+        Objects.requireNonNull(params);
+        return new MockDefinitionContext(currentStep, taskName, null, methodName, params);
+    }
+}

plugins/tasks/mock/src/main/java/com/walmartlabs/concord/plugins/mock/MockDefinitionProvider.java

@@ -21,6 +21,8 @@
  */
 
 import com.walmartlabs.concord.plugins.mock.matcher.ArgsMatcher;
+import com.walmartlabs.concord.runtime.v2.model.AbstractStep;
+import com.walmartlabs.concord.runtime.v2.runner.logging.LogUtils;
 import com.walmartlabs.concord.runtime.v2.sdk.Context;
 import com.walmartlabs.concord.runtime.v2.sdk.UserDefinedException;
 import com.walmartlabs.concord.runtime.v2.sdk.Variables;
@@ -34,18 +36,20 @@
 @Singleton
 public class MockDefinitionProvider {
 
+    private final MockDefinitionMatcher mockDefinitionMatcher = new MockDefinitionMatcher();
+
     public MockDefinition find(Context ctx, String taskName, Variables input) {
         return findMockDefinitions(ctx, mock ->
-                taskName.equals(mock.task()) && ArgsMatcher.match(input.toMap(), mock.input()));
+                mockDefinitionMatcher.matches(MockDefinitionContext.task(ctx.execution().currentStep(), taskName, input), mock));
     }
 
     public MockDefinition find(Context ctx, String taskName, String method, Object[] params) {
         return findMockDefinitions(ctx, mock ->
-                taskName.equals(mock.task()) && method.equals(mock.method()) && ArgsMatcher.match(params, mock.args()));
+                mockDefinitionMatcher.matches(MockDefinitionContext.method(ctx.execution().currentStep(), taskName, method, params), mock));
     }
 
     public boolean isTaskMocked(Context ctx, String taskName) {
-        return mocks(ctx).anyMatch(mock -> taskName.equals(mock.task()));
+        return mocks(ctx).anyMatch(mock -> ArgsMatcher.match(taskName, mock.task()));
     }
 
     private static MockDefinition findMockDefinitions(Context ctx, Predicate<MockDefinition> predicate) {
@@ -63,4 +67,108 @@ private static Stream<MockDefinition> mocks(Context ctx) {
         return ctx.variables().getList("mocks", List.of()).stream()
                 .map(m -> new MockDefinition((Map<String, Object>) m));
     }
+
+    public static class MockDefinitionMatcher {
+
+        private final List<Matcher> matchers = List.of(
+                new TaskNameMatcher(),
+                new MethodNameMatcher(),
+                new StepNameMatcher(),
+                new StepMetaMatcher(),
+                new TaskInputMatcher(),
+                new TaskArgsMatcher()
+        );
+
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            for (Matcher matcher : matchers) {
+                if (!matcher.matches(context, mock)) {
+                    return false;
+                }
+            }
+            return true;
+        }
+    }
+
+    public interface Matcher {
+
+        boolean matches(MockDefinitionContext context, MockDefinition mock);
+    }
+
+    public static class TaskNameMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            return ArgsMatcher.match(context.taskName(), mock.task());
+        }
+    }
+
+    public static class MethodNameMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            if (context.method() == null) {
+                return true;
+            }
+
+            return ArgsMatcher.match(context.method(), mock.method());
+        }
+    }
+
+    public static class StepNameMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            if (mock.stepName() == null) {
+                return true;
+            }
+
+            var logContext = LogUtils.getContext();
+            return logContext != null && ArgsMatcher.match(logContext.segmentName(), mock.stepName());
+        }
+    }
+
+    public static class StepMetaMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            if (mock.stepMeta().isEmpty()) {
+                return true;
+            }
+
+            if (!(context.currentStep() instanceof AbstractStep<?> step)) {
+                return false;
+            }
+
+            var stepOptions = step.getOptions();
+            if (stepOptions == null) {
+                return false;
+            }
+
+            return ArgsMatcher.match(stepOptions.meta(), mock.stepMeta());
+        }
+    }
+
+    public static class TaskInputMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            if (context.input() == null) {
+                return true;
+            }
+
+            return ArgsMatcher.match(context.input().toMap(), mock.input());
+        }
+    }
+
+    public static class TaskArgsMatcher implements Matcher {
+
+        @Override
+        public boolean matches(MockDefinitionContext context, MockDefinition mock) {
+            if (context.params() == null) {
+                return true;
+            }
+
+            return ArgsMatcher.match(context.params(), mock.args());
+        }
+    }
 }

plugins/tasks/mock/src/test/java/com/walmartlabs/concord/plugins/mock/MockDefinitionMatcherTest.java

@@ -0,0 +1,186 @@
+package com.walmartlabs.concord.plugins.mock;
+
+/*-
+ * *****
+ * Concord
+ * -----
+ * Copyright (C) 2017 - 2024 Walmart Inc.
+ * -----
+ * 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
+ * 
+ *      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.
+ * =====
+ */
+
+import com.walmartlabs.concord.runtime.v2.model.Location;
+import com.walmartlabs.concord.runtime.v2.model.Step;
+import com.walmartlabs.concord.runtime.v2.model.TaskCall;
+import com.walmartlabs.concord.runtime.v2.model.TaskCallOptions;
+import com.walmartlabs.concord.runtime.v2.sdk.MapBackedVariables;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+
+import static com.walmartlabs.concord.plugins.mock.MockDefinitionProvider.*;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+public class MockDefinitionMatcherTest {
+
+    private MockDefinitionMatcher mockDefinitionMatcher;
+
+    @BeforeEach
+    public void setUp() {
+        mockDefinitionMatcher = new MockDefinitionMatcher();
+    }
+
+    @Test
+    public void testTaskMatch() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+           param2: value2
+         */
+        var context = MockDefinitionContext.task(mock(Step.class), "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask",
+                "in", Map.of("param1", "value1")
+        ));
+
+        assertTrue(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testMatchOnlyByMeta() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+           param2: value2
+         */
+        var currentStep = new TaskCall(Location.builder().build(), "myTask", TaskCallOptions.builder().meta(Map.of("taskId", "BOO")).build());
+        var context = MockDefinitionContext.task(currentStep, "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask",
+                "stepMeta", Map.of("taskId", "BO.*")
+        ));
+
+        assertTrue(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testTaskMatchByMeta() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+         meta:
+           taskId: "BOO"
+         */
+        var currentStep = new TaskCall(Location.builder().build(), "myTask", TaskCallOptions.builder().meta(Map.of("taskId", "BOO")).build());
+        var context = MockDefinitionContext.task(currentStep, "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask",
+                "in", Map.of("param1", "value1"),
+                "stepMeta", Map.of("taskId", "BO.*")
+        ));
+
+        assertTrue(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testTaskMethodMatch() {
+        // expr: ${myTask.myMethod(1, 2)}
+        var context = MockDefinitionContext.method(mock(Step.class), "myTask", "myMethod", new Object[] {1, 2});
+
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask",
+                "method", "myMethod",
+                "args", List.of(1, 2)
+        ));
+
+        assertTrue(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testTaskMethodMatchByMeta() {
+        // expr: ${myTask.myMethod(1, 2)}
+        // meta:
+        //   taskId: "BOO"
+        var currentStep = new TaskCall(Location.builder().build(), "myTask", TaskCallOptions.builder().meta(Map.of("taskId", "BOO")).build());
+        var context = MockDefinitionContext.method(currentStep, "myTask", "myMethod", new Object[] {1, 2});
+
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask",
+                "method", "myMethod",
+                "args", List.of(1, 2),
+                "stepMeta", Map.of("taskId", "BO.*")
+        ));
+
+        assertTrue(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testNotMatch_taskName() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+           param2: value2
+         */
+        var context = MockDefinitionContext.task(mock(Step.class), "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask2",
+                "in", Map.of("param1", "value1")
+        ));
+
+        // real taskName=myTask, mocked: "myTask2"
+        assertFalse(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testNotMatch_InputParams() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+           param2: value2
+         */
+        var context = MockDefinitionContext.task(mock(Step.class), "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask2",
+                "in", Map.of("param3", "value3")
+        ));
+
+        assertFalse(mockDefinitionMatcher.matches(context, mock));
+    }
+
+    @Test
+    public void testNotMatch_Meta() {
+        /*
+         task: myTask
+         in:
+           param1: value1
+           param2: value2
+         */
+        var context = MockDefinitionContext.task(mock(Step.class), "myTask", new MapBackedVariables(Map.of("param1", "value1", "param2", "value2")));
+        var mock = new MockDefinition(Map.of(
+                "task", "myTask2",
+                "stepMeta", Map.of("taskId", "BO.*")
+        ));
+
+        assertFalse(mockDefinitionMatcher.matches(context, mock));
+    }
+}