/digdag-core/src/test/java/io/digdag/core/workflow/OperatorTestingUtils.java

https://github.com/treasure-data/digdag · Java · 173 lines · 151 code · 22 blank · 0 comment · 0 complexity · a4ec8ba84967283c4fbbed983a2ba53b MD5 · raw file

  1. package io.digdag.core.workflow;
  2. import com.google.common.base.Optional;
  3. import com.google.common.collect.ImmutableMap;
  4. import com.google.inject.Guice;
  5. import com.google.inject.Injector;
  6. import com.google.inject.Module;
  7. import com.google.inject.Scopes;
  8. import com.google.inject.TypeLiteral;
  9. import io.digdag.client.config.Config;
  10. import io.digdag.client.config.ConfigFactory;
  11. import io.digdag.client.config.ConfigUtils;
  12. import io.digdag.core.Environment;
  13. import io.digdag.core.agent.ConfigEvalEngine;
  14. import io.digdag.core.agent.GrantedPrivilegedVariables;
  15. import io.digdag.core.agent.OperatorRegistry;
  16. import io.digdag.core.agent.TaskContextCommandLogger;
  17. import io.digdag.spi.CommandExecutor;
  18. import io.digdag.spi.CommandLogger;
  19. import io.digdag.spi.ImmutableTaskRequest;
  20. import io.digdag.spi.OperatorContext;
  21. import io.digdag.spi.OperatorFactory;
  22. import io.digdag.spi.PrivilegedVariables;
  23. import io.digdag.spi.SecretProvider;
  24. import io.digdag.spi.TaskRequest;
  25. import io.digdag.spi.TemplateEngine;
  26. import java.nio.file.Path;
  27. import java.util.Properties;
  28. import java.util.Map;
  29. import java.util.UUID;
  30. import java.time.Instant;
  31. import java.time.ZoneId;
  32. import static io.digdag.client.config.ConfigUtils.newConfig;
  33. public class OperatorTestingUtils
  34. {
  35. private OperatorTestingUtils()
  36. { }
  37. public static <T extends OperatorFactory> T newOperatorFactory(Class<T> factoryClass)
  38. {
  39. Injector initInjector = Guice.createInjector((Module) (binder) -> {
  40. binder.bind(CommandExecutor.class).to(MockCommandExecutor.class).in(Scopes.SINGLETON);
  41. binder.bind(CommandLogger.class).to(TaskContextCommandLogger.class).in(Scopes.SINGLETON);
  42. binder.bind(TemplateEngine.class).to(ConfigEvalEngine.class).in(Scopes.SINGLETON);
  43. binder.bind(ConfigFactory.class).toInstance(ConfigUtils.configFactory);
  44. binder.bind(Config.class).toInstance(newConfig());
  45. binder.bind(OperatorRegistry.DynamicOperatorPluginInjectionModule.class).in(Scopes.SINGLETON);
  46. binder.bind(new TypeLiteral<Map<String, String>>() {}).annotatedWith(Environment.class).toInstance(ImmutableMap.<String, String>of());
  47. });
  48. Injector injector = Guice.createInjector(
  49. (Module) initInjector.getInstance(OperatorRegistry.DynamicOperatorPluginInjectionModule.class),
  50. (Module) (binder) -> binder.bind(factoryClass)
  51. );
  52. return injector.getInstance(factoryClass);
  53. }
  54. public static ImmutableTaskRequest newTaskRequest()
  55. {
  56. return ImmutableTaskRequest.builder()
  57. .siteId(1)
  58. .projectId(2)
  59. .workflowName("wf")
  60. .revision(Optional.of("rev"))
  61. .taskId(3)
  62. .attemptId(4)
  63. .sessionId(5)
  64. .taskName("t")
  65. .lockId("l")
  66. .timeZone(ZoneId.systemDefault())
  67. .sessionUuid(UUID.randomUUID())
  68. .sessionTime(Instant.now())
  69. .createdAt(Instant.now())
  70. .config(newConfig())
  71. .localConfig(newConfig())
  72. .lastStateParams(newConfig())
  73. .build();
  74. }
  75. public static TestingOperatorContext newContext(Path projectPath, TaskRequest request)
  76. {
  77. return new TestingOperatorContext(
  78. projectPath,
  79. request,
  80. TestingSecretProvider.empty(),
  81. GrantedPrivilegedVariables.empty());
  82. }
  83. public static TestingOperatorContext newContext(Path projectPath, TaskRequest request, SecretProvider secrets)
  84. {
  85. return new TestingOperatorContext(
  86. projectPath,
  87. request,
  88. secrets,
  89. GrantedPrivilegedVariables.empty());
  90. }
  91. public static class TestingOperatorContext
  92. implements OperatorContext
  93. {
  94. private final Path projectPath;
  95. private final TaskRequest taskRequest;
  96. private final SecretProvider secrets;
  97. private final PrivilegedVariables privilegedVariables;
  98. public TestingOperatorContext(
  99. Path projectPath,
  100. TaskRequest taskRequest,
  101. SecretProvider secrets,
  102. PrivilegedVariables privilegedVariables)
  103. {
  104. this.projectPath = projectPath;
  105. this.taskRequest = taskRequest;
  106. this.secrets = secrets;
  107. this.privilegedVariables = privilegedVariables;
  108. }
  109. @Override
  110. public Path getProjectPath()
  111. {
  112. return projectPath;
  113. }
  114. @Override
  115. public TaskRequest getTaskRequest()
  116. {
  117. return taskRequest;
  118. }
  119. @Override
  120. public SecretProvider getSecrets()
  121. {
  122. return secrets;
  123. }
  124. @Override
  125. public PrivilegedVariables getPrivilegedVariables()
  126. {
  127. return privilegedVariables;
  128. }
  129. }
  130. private static class TestingSecretProvider
  131. implements SecretProvider
  132. {
  133. public static TestingSecretProvider empty()
  134. {
  135. return fromProperties(new Properties());
  136. }
  137. public static TestingSecretProvider fromProperties(Properties props)
  138. {
  139. return new TestingSecretProvider(props);
  140. }
  141. private final Properties props;
  142. private TestingSecretProvider(Properties props)
  143. {
  144. this.props = props;
  145. }
  146. public Optional<String> getSecretOptional(String key)
  147. {
  148. return Optional.fromNullable(props.getProperty(key));
  149. }
  150. }
  151. }