Lab 5.3: Build It — MiniTezCluster Integration Test

Lab type: Build It — a full mini-cluster integration test from scratch Estimated time: 150 min Tez module: tez-tests Key class (you create): org.apache.tez.test.TestMiniClusterOrderedWordCount

Background

In Lab 5.1 you read TestTezJobs. Now you write your own integration test in the same tree, using the same MiniDFSCluster + MiniTezCluster boilerplate, but you drive the DAG through the client API (TezClient.submitDAG → DAGClient → DAGStatus) so you can assert on the final DAG state and on TezCounters — not just "the job returned 0."

This is the exact artifact a committer expects to accompany a DAG-level change: a test that submits a real multi-vertex DAG on the mini-cluster, waits for completion, asserts DAGStatus.State.SUCCEEDED, and verifies a specific, derivable counter value. You will build it against the real OrderedWordCount example DAG (Tokenizer → Summation → Sorter) so every class and counter is real and every command runs.

Note on "using the umbilical properly": In a mini-cluster test you do not touch the task↔AM umbilical (TaskCommunicator/TezTaskUmbilicalProtocol) directly — that is wired automatically when the AM launches real containers. What you do control is the client-side protocol: TezClient → AM submission, and DAGClient polling for status and counters. Manipulating the umbilical by hand is a unit-test concern (see TestUmbilical in the runtime tests); at the integration tier you let the real umbilical carry heartbeats and events, and you assert on the results it produces.

Why This Lab Matters for Contributors

• "Add a DAG feature" or "fix a DAG-level bug" patches are only merged with an integration test that asserts the observable outcome — final state plus a counter. This lab produces exactly that.
• Driving TezClient/DAGClient yourself (rather than the Tool.run shortcut) teaches the client protocol you will need to debug submission and status bugs.
• Deriving the expected counter by hand — instead of asserting whatever came out — is the discipline that separates a test that proves something from one that merely runs.

Prerequisites

• Labs 5.1 and 5.2 complete; export TEZ_SRC=/path/to/tez; mvn install -DskipTests -q succeeded.
• You can run TestTezJobs#testOrderedWordCount to green.
• Read Counters & Diagnostics and DAG Client.

Confirm the APIs you will call are real before you write against them:

cd "$TEZ_SRC"
grep -n "public synchronized DAGClient submitDAG" tez-api/src/main/java/org/apache/tez/client/TezClient.java
grep -n "waitForCompletionWithStatusUpdates\|getVertexStatus" tez-api/src/main/java/org/apache/tez/dag/api/client/DAGClient.java
grep -n "public enum State\|getState\|getDAGCounters" tez-api/src/main/java/org/apache/tez/dag/api/client/DAGStatus.java
grep -n "createDAG" tez-examples/src/main/java/org/apache/tez/examples/OrderedWordCount.java

Step-by-Step Tasks

Step 1: Understand the DAG you will submit

OrderedWordCount.createDAG(...) builds a 3-vertex DAG (tez-examples):

Tokenizer  --(OrderedPartitionedKV edge)-->  Summation  --(OrderedPartitionedKV edge)-->  Sorter
 (reads MRInput text,                         (sums counts,                                 (no-op,
  emits <word, 1>)                             emits <count, word>)                          writes MROutput)

Its signature (verify with the grep above):

public static DAG createDAG(TezConfiguration tezConf, String inputPath, String outputPath,
    int numPartitions, boolean disableSplitGrouping, boolean isGenerateSplitInClient, String dagName)

The vertex names are the string constants "Tokenizer", "Summation", "Sorter" (the sink on Sorter is named "Output"). You will read a per-vertex counter off the "Tokenizer" vertex.

Step 2: Derive the counter you will assert

Reuse the standard input generator (identical to TestTezJobs.generateOrderedWordCountInput): for each i in 1..10, write the word a_i exactly (11 - i) times to each of two files. Total word occurrences:

sum over i=1..10 of (11 - i)  =  10 + 9 + ... + 1  =  55  per file
55 * 2 files                  =  110 total tokens

The Tokenizer emits one <word, 1> record per token, so the Tokenizer vertex's TaskCounter.OUTPUT_RECORDS must equal 110. That is your hand-derived, deterministic assertion — not a value you copied from a passing run.

Tip: Deriving the number first is the whole point. If your test just asserts "whatever the run produced," it can't catch a regression that changes the count. A committer will ask "how do you know 110 is right?" — and "I did the arithmetic" is the only acceptable answer.

Step 3: Create the test file

Create tez-tests/src/test/java/org/apache/tez/test/TestMiniClusterOrderedWordCount.java. No pom.xml changes are needed — tez-tests already depends on tez-examples, hadoop-minicluster, and JUnit. Here is the complete, runnable test:

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */
package org.apache.tez.test;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import java.io.IOException;
import java.util.EnumSet;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hdfs.MiniDFSCluster;
import org.apache.tez.client.TezClient;
import org.apache.tez.common.counters.TaskCounter;
import org.apache.tez.common.counters.TezCounter;
import org.apache.tez.common.counters.TezCounters;
import org.apache.tez.dag.api.DAG;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.api.client.DAGClient;
import org.apache.tez.dag.api.client.DAGStatus;
import org.apache.tez.dag.api.client.StatusGetOpts;
import org.apache.tez.dag.api.client.VertexStatus;
import org.apache.tez.examples.OrderedWordCount;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;

public class TestMiniClusterOrderedWordCount {

  private static MiniTezCluster tezCluster;
  private static MiniDFSCluster dfsCluster;
  private static Configuration conf = new Configuration();
  private static FileSystem remoteFs;

  private static final String TEST_ROOT_DIR =
      "target" + Path.SEPARATOR + TestMiniClusterOrderedWordCount.class.getName() + "-tmpDir";

  @BeforeClass
  public static void setup() throws IOException {
    // 1) Real in-JVM HDFS.
    conf.set(MiniDFSCluster.HDFS_MINIDFS_BASEDIR, TEST_ROOT_DIR);
    dfsCluster = new MiniDFSCluster.Builder(conf).numDataNodes(2).format(true).racks(null).build();
    remoteFs = dfsCluster.getFileSystem();

    // 2) Real in-JVM YARN + Tez AM. Must come AFTER remoteFs exists so fs.defaultFS points at HDFS.
    tezCluster = new MiniTezCluster(TestMiniClusterOrderedWordCount.class.getName(), 1, 1, 1);
    conf.set("fs.defaultFS", remoteFs.getUri().toString());
    conf.setLong(TezConfiguration.TEZ_AM_SLEEP_TIME_BEFORE_EXIT_MILLIS, 500);
    tezCluster.init(conf);
    tezCluster.start();
  }

  @AfterClass
  public static void tearDown() {
    if (tezCluster != null) {
      tezCluster.stop();
      tezCluster = null;
    }
    if (dfsCluster != null) {
      dfsCluster.shutdown();
      dfsCluster = null;
    }
  }

  /** Same input the canonical TestTezJobs uses: a_i written (11-i) times to EACH of two files. */
  private static void generateInput(Path inputDir, FileSystem fs) throws IOException {
    FSDataOutputStream f1 = fs.create(new Path(inputDir, "inPath1"));
    FSDataOutputStream f2 = fs.create(new Path(inputDir, "inPath2"));
    try {
      for (int i = 1; i <= 10; ++i) {
        String word = "a_" + i;
        for (int j = 10; j >= i; --j) {   // (11 - i) copies
          f1.write(word.getBytes());
          f1.writeChars("\t");
          f2.write(word.getBytes());
          f2.writeChars("\t");
        }
      }
      f1.hsync();
      f2.hsync();
    } finally {
      f1.close();
      f2.close();
    }
  }

  @Test(timeout = 120000)
  public void testOrderedWordCountCountersAndState() throws Exception {
    Path inputDir = new Path("/tmp/owc-mini-input");
    Path outputDir = new Path("/tmp/owc-mini-output");
    Path stagingDir = new Path("/tmp/owc-mini-staging");
    remoteFs.mkdirs(inputDir);
    remoteFs.mkdirs(stagingDir);
    remoteFs.delete(outputDir, true);          // MROutput requires a non-existent output dir
    generateInput(inputDir, remoteFs);

    TezConfiguration tezConf = new TezConfiguration(tezCluster.getConfig());
    tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDir.toString());

    TezClient tezClient = TezClient.create("OwcMiniClusterTest", tezConf, true /* session */);
    tezClient.start();
    try {
      tezClient.waitTillReady();

      DAG dag = OrderedWordCount.createDAG(tezConf, inputDir.toString(), outputDir.toString(),
          2 /* numPartitions */, false /* disableSplitGrouping */,
          false /* generateSplitsInClient */, "OwcMiniClusterTest");

      DAGClient dagClient = tezClient.submitDAG(dag);
      DAGStatus dagStatus =
          dagClient.waitForCompletionWithStatusUpdates(EnumSet.of(StatusGetOpts.GET_COUNTERS));

      // Assertion 1: final DAG state.
      assertEquals("DAG must succeed", DAGStatus.State.SUCCEEDED, dagStatus.getState());

      // Assertion 2: DAG counters are present.
      TezCounters dagCounters = dagStatus.getDAGCounters();
      assertNotNull("DAG counters must be present", dagCounters);

      // Assertion 3: derived per-vertex counter. Tokenizer emits one record per token; the input
      // has sum(i=1..10)(11-i) = 55 tokens per file * 2 files = 110 tokens total.
      VertexStatus tokenizerStatus =
          dagClient.getVertexStatus("Tokenizer", EnumSet.of(StatusGetOpts.GET_COUNTERS));
      assertNotNull(tokenizerStatus);
      TezCounter tokenizerOutput =
          tokenizerStatus.getVertexCounters().findCounter(TaskCounter.OUTPUT_RECORDS);
      assertNotNull(tokenizerOutput);
      assertEquals("Tokenizer must emit one record per input token", 110L, tokenizerOutput.getValue());
      assertTrue("Sanity: DAG output records must be positive",
          tokenizerOutput.getValue() > 0);
    } finally {
      remoteFs.delete(stagingDir, true);
      tezClient.stop();
    }
  }
}

Step 4: Build and run

cd "$TEZ_SRC"
mvn install -DskipTests -q
mvn test -pl tez-tests -Dtest='TestMiniClusterOrderedWordCount#testOrderedWordCountCountersAndState' 2>&1 | tail -25

Expected tail:

Tests run: 1, Failures: 0, Errors: 0, Skipped: 0
BUILD SUCCESS

Step 5: Read status the way the test does

While the DAG runs, the client polls the AM for status. Watch it in the captured log:

grep -n "Submitting dag\|DAG State\|SUCCEEDED\|Tokenizer\|OUTPUT_RECORDS" \
  tez-tests/target/surefire-reports/org.apache.tez.test.TestMiniClusterOrderedWordCount-output.txt | head

Confirm you can see the DAG reach SUCCEEDED and (with GET_COUNTERS) the counters being fetched. If you drop StatusGetOpts.GET_COUNTERS, getDAGCounters() may return counters only for a finished DAG — the flag forces counters into the status you wait on. This is a common integration-test bug: asking for a counter you never requested.

Step 5b: The client polling protocol and DAG states

waitForCompletionWithStatusUpdates is a blocking poll loop: the DAGClient repeatedly asks the AM for the current DAGStatus until the DAG reaches a terminal state. The DAGStatus.State values you can observe (from DAGStatus.State) are:

You gate on SUCCEEDED. On any other terminal state, dagStatus.getDiagnostics() tells you why — that is the first thing to print when an integration test fails:

if (dagStatus.getState() != DAGStatus.State.SUCCEEDED) {
  System.err.println("DAG failed: " + dagStatus.getDiagnostics());
}

Passing EnumSet.of(StatusGetOpts.GET_COUNTERS) tells the AM to include counters in every status response it sends back during the poll — without it, getDAGCounters() and getVertexCounters() can be empty. This is the single most common "why are my counters null?" bug in a first integration test.

Step 6: Prove the counter assertion bites

Change 110L to 111L and re-run. The test must fail with:

java.lang.AssertionError: Tokenizer must emit one record per input token expected:<111> but was:<110>

Change it back. A test whose numeric assertion you have watched fail on a wrong value is a test you can trust. Revert to 110L.

Step 7: Compare the two submission styles

You submitted via TezClient.submitDAG + DAGClient. TestTezJobs#testOrderedWordCount submits via OrderedWordCount.run(...) (the Tool path). Write two or three sentences: which style lets you assert on counters, which is terser, and when you would choose each.

Implementation Requirements / Deliverables

• TestMiniClusterOrderedWordCount.java created in tez-tests/src/test/java/org/apache/tez/test/, compiling and passing.
• The @BeforeClass wires MiniDFSCluster then MiniTezCluster, with fs.defaultFS set after remoteFs exists.
• Three assertions: DAGStatus.State.SUCCEEDED; non-null getDAGCounters(); Tokenizer OUTPUT_RECORDS == 110 derived by hand (show the arithmetic in a comment).
• Evidence the counter assertion bites (Step 6: the 111L failure, then reverted).
• A short note comparing submitDAG/DAGClient vs the Tool.run path (Step 7).

Troubleshooting

Stretch Goals

1. Assert a DAG-level counter. From dagStatus.getDAGCounters(), find TaskCounter.OUTPUT_RECORDS aggregated across all vertices, derive its expected value (Tokenizer 110 + Summation 10 distinct words
   • Sorter 10), and assert it. Reason carefully about grouping.
2. Force a failure and assert diagnostics. Point the input at a non-existent path, submit, and assert dagStatus.getState() == FAILED and that dagStatus.getDiagnostics() is non-empty. This is how you test the unhappy path.
3. Parameterize split grouping. Turn the test into a JUnit @Parameterized class that runs with disableSplitGrouping both true and false, mirroring how TestTezJobs covers both.
4. Session reuse. Submit the DAG twice on the same started TezClient session and assert both succeed — proving you understand session vs per-DAG AM lifecycle.

Deeper Understanding

Validation / Self-check

1. Why must conf.set("fs.defaultFS", ...) come after remoteFs = dfsCluster.getFileSystem()?
2. Which API returns counters, and what must you pass to get them into the DAGStatus you wait on?
3. Show the arithmetic that yields 110 for the Tokenizer's OUTPUT_RECORDS.
4. What does "using the umbilical properly" mean at the integration tier, and why don't you touch it directly here?
5. What is the difference between submitting via TezClient.submitDAG/DAGClient and via the example's Tool.run? Which can assert on counters?
6. How did you prove your 110L assertion actually guards against regressions?
7. Why does MROutput require you to delete the output directory before submitting?

When your test passes, its counter assertion bites, and you can explain the client protocol, continue to Lab 5.4: Fix It — Un-Ignore a Flaky Test.