test

Will-Lo · Jan 18, 2024 · 0bc3373 · 0bc3373
1 parent 318c661
commit 0bc3373
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 16 deletions.
diff --git a/gobblin-service/src/test/java/org/apache/gobblin/service/GobblinServiceManagerTest.java b/gobblin-service/src/test/java/org/apache/gobblin/service/GobblinServiceManagerTest.java
@@ -28,6 +28,9 @@
 
 import org.apache.commons.io.FileUtils;
 import org.apache.curator.test.TestingServer;
+
+import org.apache.gobblin.service.modules.orchestration.AbstractUserQuotaManager;
+import org.apache.gobblin.service.modules.orchestration.ServiceAzkabanConfigKeys;
 import org.apache.hadoop.fs.Path;
 import org.eclipse.jetty.http.HttpStatus;
 import org.eclipse.jgit.api.Git;
@@ -36,7 +39,6 @@
 import org.eclipse.jgit.util.FS;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.testcontainers.containers.MySQLContainer;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -60,9 +62,7 @@
 import org.apache.gobblin.runtime.spec_catalog.FlowCatalog;
 import org.apache.gobblin.service.modules.core.GobblinServiceManager;
 import org.apache.gobblin.service.modules.flow.MockedSpecCompiler;
-import org.apache.gobblin.service.modules.orchestration.AbstractUserQuotaManager;
 import org.apache.gobblin.service.modules.orchestration.DagManager;
-import org.apache.gobblin.service.modules.orchestration.ServiceAzkabanConfigKeys;
 import org.apache.gobblin.service.monitoring.FsJobStatusRetriever;
 import org.apache.gobblin.service.monitoring.GitConfigMonitor;
 import org.apache.gobblin.testing.AssertWithBackoff;
@@ -120,7 +120,10 @@ public class GobblinServiceManagerTest {
   private GobblinServiceManager gobblinServiceManager;
   private FlowConfigV2Client flowConfigClient;
 
-  private MySQLContainer mysql;
+  private ITestMetastoreDatabase mysql;
+
+  private static final String TEST_USER = "testUser";
+  private static final String TEST_PASSWORD = "testPassword";
 
   private Git gitForPush;
   private TestingServer testingServer;
@@ -136,23 +139,17 @@ public void setup() throws Exception {
     cleanUpDir(SERVICE_WORK_DIR);
     cleanUpDir(SPEC_STORE_PARENT_DIR);
 
-    mysql = new MySQLContainer("mysql:" + TestServiceDatabaseConfig.MysqlVersion);
-    mysql.start();
-    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_URL_KEY, mysql.getJdbcUrl());
-    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_USERNAME, mysql.getUsername());
-    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_PASSWORD, mysql.getPassword());
-
-    ITestMetastoreDatabase testMetastoreDatabase = TestMetastoreDatabaseFactory.get();
+    mysql = TestMetastoreDatabaseFactory.get();
 
     testingServer = new TestingServer(true);
 
     flowProperties.put("param1", "value1");
     flowProperties.put(ServiceConfigKeys.FLOW_SOURCE_IDENTIFIER_KEY, TEST_SOURCE_NAME);
     flowProperties.put(ServiceConfigKeys.FLOW_DESTINATION_IDENTIFIER_KEY, TEST_SINK_NAME);
 
-    serviceCoreProperties.put(ConfigurationKeys.STATE_STORE_DB_USER_KEY, "testUser");
-    serviceCoreProperties.put(ConfigurationKeys.STATE_STORE_DB_PASSWORD_KEY, "testPassword");
-    serviceCoreProperties.put(ConfigurationKeys.STATE_STORE_DB_URL_KEY, testMetastoreDatabase.getJdbcUrl());
+    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_URL_KEY, mysql.getJdbcUrl());
+    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_USERNAME, TEST_USER);
+    serviceCoreProperties.put(ServiceConfigKeys.SERVICE_DB_PASSWORD, TEST_PASSWORD);
     serviceCoreProperties.put("zookeeper.connect", testingServer.getConnectString());
     serviceCoreProperties.put(ConfigurationKeys.STATE_STORE_FACTORY_CLASS_KEY, MysqlJobStatusStateStoreFactory.class.getName());
 
@@ -259,8 +256,6 @@ public void cleanUp() throws Exception {
     } catch(Exception e) {
       System.err.println("Failed to close ZK testing server.");
     }
-
-    mysql.stop();
   }
 
   /**

diff --git a/...main/java/org/apache/gobblin/temporal/ddm/workflow/impl/ProcessWorkUnitsWorkflowImpl.java b/...main/java/org/apache/gobblin/temporal/ddm/workflow/impl/ProcessWorkUnitsWorkflowImpl.java
@@ -36,6 +36,7 @@
 import org.apache.gobblin.temporal.util.nesting.work.WorkflowAddr;
 import org.apache.gobblin.temporal.util.nesting.work.Workload;
 import org.apache.gobblin.temporal.util.nesting.workflow.NestingExecWorkflow;
+import org.slf4j.MDC;
 
 
 @Slf4j
@@ -45,6 +46,7 @@ public class ProcessWorkUnitsWorkflowImpl implements ProcessWorkUnitsWorkflow {
 
   @Override
   public int process(WUProcessingSpec workSpec) {
+    MDC.put("workflowId", Workflow.getInfo().getWorkflowId());
     Workload<WorkUnitClaimCheck> workload = createWorkload(workSpec);
     NestingExecWorkflow<WorkUnitClaimCheck> processingWorkflow = createProcessingWorkflow(workSpec);
     int workunitsProcessed = processingWorkflow.performWorkload(

diff --git a/...va/org/apache/gobblin/temporal/util/nesting/workflow/AbstractNestingExecWorkflowImpl.java b/...va/org/apache/gobblin/temporal/util/nesting/workflow/AbstractNestingExecWorkflowImpl.java
@@ -36,6 +36,7 @@
 
 import org.apache.gobblin.temporal.util.nesting.work.WorkflowAddr;
 import org.apache.gobblin.temporal.util.nesting.work.Workload;
+import org.slf4j.MDC;
 
 
 /** Core skeleton of {@link NestingExecWorkflow}: realizing classes need only define {@link #launchAsyncActivity} */