Merge LogDnsClient's QueryLeafIndex and QueryAuditProof methods

components/certificate_transparency/log_dns_client_unittest.cc

Index: components/certificate_transparency/log_dns_client_unittest.cc
diff --git a/components/certificate_transparency/log_dns_client_unittest.cc b/components/certificate_transparency/log_dns_client_unittest.cc
index c1e914297fc8e207edea5522e0e4c6c614473596..b41c06169de72ec384a982d0666d7eeae78c1be0 100644
--- a/components/certificate_transparency/log_dns_client_unittest.cc
+++ b/components/certificate_transparency/log_dns_client_unittest.cc
@@ -39,27 +39,29 @@ using ::testing::Le;
 using ::testing::Not;
 using ::testing::NotNull;
 using net::test::IsError;
 using net::test::IsOk;
 
+// Sample Merkle leaf hashes.
 const char* kLeafHashes[] = {
     "\x1f\x25\xe1\xca\xba\x4f\xf9\xb8\x27\x24\x83\x0f\xca\x60\xe4\xc2\xbe\xa8"
     "\xc3\xa9\x44\x1c\x27\xb0\xb4\x3e\x6a\x96\x94\xc7\xb8\x04",
     "\x2c\x26\xb4\x6b\x68\xff\xc6\x8f\xf9\x9b\x45\x3c\x1d\x30\x41\x34\x13\x42"
     "\x2d\x70\x64\x83\xbf\xa0\xf9\x8a\x5e\x88\x62\x66\xe7\xae",
     "\xfc\xde\x2b\x2e\xdb\xa5\x6b\xf4\x08\x60\x1f\xb7\x21\xfe\x9b\x5c\x33\x8d"
     "\x10\xee\x42\x9e\xa0\x4f\xae\x55\x11\xb6\x8f\xbf\x8f\xb9",
 };
 
-// Assumes log domain is "ct.test"
-const char* kBase32LeafHashes[] = {
+// DNS query names for looking up the leaf index associated with each hash in
+// |kLeafHashes|. Assumes the log domain is "ct.test".
+const char* kLeafIndexQnames[] = {
     "D4S6DSV2J743QJZEQMH4UYHEYK7KRQ5JIQOCPMFUHZVJNFGHXACA.hash.ct.test.",
     "FQTLI23I77DI76M3IU6B2MCBGQJUELLQMSB37IHZRJPIQYTG46XA.hash.ct.test.",
     "7TPCWLW3UVV7ICDAD63SD7U3LQZY2EHOIKPKAT5OKUI3ND57R64Q.hash.ct.test.",
 };
 
-// Leaf indices and tree sizes for use with above leaf hashes.
+// Leaf indices and tree sizes for use with |kLeafHashes|.
 const uint64_t kLeafIndices[] = {0, 1, 2};
 const uint64_t kTreeSizes[] = {100, 10000, 1000000};
 
 // Only 7 audit proof nodes can fit into a DNS response, because they are sent
 // in a TXT RDATA string, which has a maximum size of 255 bytes, and each node
@@ -82,39 +84,10 @@ std::vector<std::string> GetSampleAuditProof(size_t length) {
   }
 
   return audit_proof;
 }
 
-class MockLeafIndexCallback {
- public:
-  MockLeafIndexCallback() : called_(false) {}
-
-  bool called() const { return called_; }
-  int net_error() const { return net_error_; }
-  uint64_t leaf_index() const { return leaf_index_; }
-
-  void Run(int net_error, uint64_t leaf_index) {
-    EXPECT_TRUE(!called_) << "Callback invoked more than once";
-    called_ = true;
-    net_error_ = net_error;
-    leaf_index_ = leaf_index;
-    run_loop_.Quit();
-  }
-
-  LogDnsClient::LeafIndexCallback AsCallback() {
-    return base::Bind(&MockLeafIndexCallback::Run, base::Unretained(this));
-  }
-
-  void WaitUntilRun() { run_loop_.Run(); }
-
- private:
-  bool called_;
-  int net_error_;
-  uint64_t leaf_index_;
-  base::RunLoop run_loop_;
-};
-
 class MockAuditProofCallback {
  public:
   MockAuditProofCallback() : called_(false) {}
 
   bool called() const { return called_; }
@@ -149,48 +122,30 @@ class LogDnsClientTest : public ::testing::TestWithParam<net::IoMode> {
     mock_dns_.SetSocketReadMode(GetParam());
     mock_dns_.InitializeDnsConfig();
   }
 
   std::unique_ptr<LogDnsClient> CreateLogDnsClient(
-      size_t max_concurrent_queries) {
-    return base::MakeUnique<LogDnsClient>(mock_dns_.CreateDnsClient(),
-                                          net::BoundNetLog(),
-                                          max_concurrent_queries);
-  }
-
-  void QueryLeafIndexAsync(LogDnsClient* log_client,
-                           base::StringPiece log_domain,
-                           const char leaf_hash[crypto::kSHA256Length],
-                           const LogDnsClient::LeafIndexCallback& callback) {
-    log_client->QueryLeafIndex(log_domain, leaf_hash, callback);
-  }
-
-  // Convenience function for calling QueryLeafIndexAsync synchronously.
-  void QueryLeafIndex(base::StringPiece log_domain,
-                      const char leaf_hash[crypto::kSHA256Length],
-                      MockLeafIndexCallback* callback) {
-    std::unique_ptr<LogDnsClient> log_client = CreateLogDnsClient(0);
-    QueryLeafIndexAsync(log_client.get(), log_domain, leaf_hash,
-                        callback->AsCallback());
-    callback->WaitUntilRun();
+      size_t kMaxConcurrentQueries) {
+    return base::MakeUnique<LogDnsClient>(
+        mock_dns_.CreateDnsClient(), net::BoundNetLog(), kMaxConcurrentQueries);
   }
 
   void QueryAuditProofAsync(LogDnsClient* log_client,
                             base::StringPiece log_domain,
-                            uint64_t leaf_index,
+                            const char leaf_hash[crypto::kSHA256Length],
                             uint64_t tree_size,
                             const LogDnsClient::AuditProofCallback& callback) {
-    log_client->QueryAuditProof(log_domain, leaf_index, tree_size, callback);
+    log_client->QueryAuditProof(log_domain, leaf_hash, tree_size, callback);
   }
 
   // Convenience function for calling QueryAuditProofAsync synchronously.
   void QueryAuditProof(base::StringPiece log_domain,
-                       uint64_t leaf_index,
+                       const char leaf_hash[crypto::kSHA256Length],
                        uint64_t tree_size,
                        MockAuditProofCallback* callback) {
     std::unique_ptr<LogDnsClient> log_client = CreateLogDnsClient(0);
-    QueryAuditProofAsync(log_client.get(), log_domain, leaf_index, tree_size,
+    QueryAuditProofAsync(log_client.get(), log_domain, leaf_hash, tree_size,
                          callback->AsCallback());
     callback->WaitUntilRun();
   }
 
   // This will be the NetworkChangeNotifier singleton for the duration of the
@@ -201,174 +156,169 @@ class LogDnsClientTest : public ::testing::TestWithParam<net::IoMode> {
   base::MessageLoopForIO message_loop_;
   // Allows mock DNS sockets to be setup.
   MockLogDnsTraffic mock_dns_;
 };
 
-TEST_P(LogDnsClientTest, QueryLeafIndex) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "123456");
-
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsOk());
-  EXPECT_THAT(callback.leaf_index(), Eq(123456u));
-}
-
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsThatLogDomainDoesNotExist) {
-  mock_dns_.ExpectRequestAndErrorResponse(kBase32LeafHashes[0],
+TEST_P(LogDnsClientTest, QueryAuditProofReportsThatLogDomainDoesNotExist) {
+  mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeNXDOMAIN);
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_NAME_NOT_RESOLVED));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsServerFailure) {
-  mock_dns_.ExpectRequestAndErrorResponse(kBase32LeafHashes[0],
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsServerFailuresDuringLeafIndexRequests) {
+  mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeSERVFAIL);
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsServerRefusal) {
-  mock_dns_.ExpectRequestAndErrorResponse(kBase32LeafHashes[0],
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsServerRefusalsDuringLeafIndexRequests) {
+  mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeREFUSED);
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryLeafIndexReportsMalformedResponseIfLeafIndexIsNotNumeric) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "foo");
+       QueryAuditProofReportsMalformedResponseIfLeafIndexIsNotNumeric) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "foo");
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryLeafIndexReportsMalformedResponseIfLeafIndexIsFloatingPoint) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "123456.0");
+       QueryAuditProofReportsMalformedResponseIfLeafIndexIsFloatingPoint) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456.0");
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryLeafIndexReportsMalformedResponseIfLeafIndexIsEmpty) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "");
+       QueryAuditProofReportsMalformedResponseIfLeafIndexIsEmpty) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "");
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryLeafIndexReportsMalformedResponseIfLeafIndexHasNonNumericPrefix) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0],
-                                              "foo123456");
+       QueryAuditProofReportsMalformedResponseIfLeafIndexHasNonNumericPrefix) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "foo123456");
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryLeafIndexReportsMalformedResponseIfLeafIndexHasNonNumericSuffix) {
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0],
-                                              "123456foo");
+       QueryAuditProofReportsMalformedResponseIfLeafIndexHasNonNumericSuffix) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456foo");
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsInvalidArgIfLogDomainIsEmpty) {
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("", kLeafHashes[0], &callback);
+TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLogDomainIsEmpty) {
+  MockAuditProofCallback callback;
+  QueryAuditProof("", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsInvalidArgIfLogDomainIsNull) {
-  MockLeafIndexCallback callback;
-  QueryLeafIndex(nullptr, kLeafHashes[0], &callback);
+TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLogDomainIsNull) {
+  MockAuditProofCallback callback;
+  QueryAuditProof(nullptr, kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsInvalidArgIfLeafHashIsInvalid) {
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", "foo", &callback);
+TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsInvalid) {
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", "foo", kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsInvalidArgIfLeafHashIsEmpty) {
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", "", &callback);
+TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsEmpty) {
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", "", kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsInvalidArgIfLeafHashIsNull) {
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", nullptr, &callback);
+TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsNull) {
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", nullptr, kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsSocketError) {
-  mock_dns_.ExpectRequestAndSocketError(kBase32LeafHashes[0],
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsSocketErrorsDuringLeafIndexRequests) {
+  mock_dns_.ExpectRequestAndSocketError(kLeafIndexQnames[0],
                                         net::ERR_CONNECTION_REFUSED);
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_CONNECTION_REFUSED));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryLeafIndexReportsTimeout) {
-  mock_dns_.ExpectRequestAndTimeout(kBase32LeafHashes[0]);
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsTimeoutsDuringLeafIndexRequests) {
+  mock_dns_.ExpectRequestAndTimeout(kLeafIndexQnames[0]);
 
-  MockLeafIndexCallback callback;
-  QueryLeafIndex("ct.test", kLeafHashes[0], &callback);
+  MockAuditProofCallback callback;
+  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_TIMED_OUT));
-  EXPECT_THAT(callback.leaf_index(), Eq(0u));
+  EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProof) {
   const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
 
+  // Expect a leaf index query first, to map the leaf hash to a leaf index.
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
+
   // It takes a number of DNS requests to retrieve the entire |audit_proof|
   // (see |kMaxProofNodesPerDnsResponse|).
   for (size_t nodes_begin = 0; nodes_begin < audit_proof.size();
        nodes_begin += kMaxProofNodesPerDnsResponse) {
     const size_t nodes_end = std::min(
@@ -378,11 +328,11 @@ TEST_P(LogDnsClientTest, QueryAuditProof) {
         base::StringPrintf("%zu.123456.999999.tree.ct.test.", nodes_begin),
         audit_proof.begin() + nodes_begin, audit_proof.begin() + nodes_end);
   }
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsOk());
   ASSERT_THAT(callback.proof(), NotNull());
   EXPECT_THAT(callback.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
@@ -391,10 +341,13 @@ TEST_P(LogDnsClientTest, QueryAuditProof) {
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofHandlesResponsesWithShortAuditPaths) {
   const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
 
+  // Expect a leaf index query first, to map the leaf hash to a leaf index.
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
+
   // Make some of the responses contain fewer proof nodes than they can hold.
   mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
                                                audit_proof.begin(),
                                                audit_proof.begin() + 1);
   mock_dns_.ExpectAuditProofRequestAndResponse("1.123456.999999.tree.ct.test.",
@@ -412,145 +365,146 @@ TEST_P(LogDnsClientTest, QueryAuditProofHandlesResponsesWithShortAuditPaths) {
   mock_dns_.ExpectAuditProofRequestAndResponse("13.123456.999999.tree.ct.test.",
                                                audit_proof.begin() + 13,
                                                audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsOk());
   ASSERT_THAT(callback.proof(), NotNull());
   EXPECT_THAT(callback.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
   // EXPECT_THAT(callback.proof()->tree_size, Eq(999999));
   EXPECT_THAT(callback.proof()->nodes, Eq(audit_proof));
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsThatLogDomainDoesNotExist) {
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsThatAuditProofQnameDoesNotExist) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeNXDOMAIN);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_NAME_NOT_RESOLVED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsServerFailure) {
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsServerFailureDuringAuditProofRequests) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeSERVFAIL);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsServerRefusal) {
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsServerRefusalDuringAuditProofRequests) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeREFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsResponseMalformedIfNodeTooShort) {
   // node is shorter than a SHA-256 hash (31 vs 32 bytes)
   const std::vector<std::string> audit_proof(1, std::string(31, 'a'));
 
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsResponseMalformedIfNodeTooLong) {
   // node is longer than a SHA-256 hash (33 vs 32 bytes)
   const std::vector<std::string> audit_proof(1, std::string(33, 'a'));
 
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsResponseMalformedIfEmpty) {
   const std::vector<std::string> audit_proof;
 
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLogDomainIsEmpty) {
-  MockAuditProofCallback callback;
-  QueryAuditProof("", 123456, 999999, &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
-}
-
-TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLogDomainIsNull) {
-  MockAuditProofCallback callback;
-  QueryAuditProof(nullptr, 123456, 999999, &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
-}
-
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsInvalidArgIfLeafIndexEqualToTreeSize) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
+
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 123456, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsInvalidArgIfLeafIndexGreaterThanTreeSize) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "999999");
+
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 999999, 123456, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsSocketError) {
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsSocketErrorsDuringAuditProofRequests) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndSocketError("0.123456.999999.tree.ct.test.",
                                         net::ERR_CONNECTION_REFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_CONNECTION_REFUSED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, QueryAuditProofReportsTimeout) {
+TEST_P(LogDnsClientTest,
+       QueryAuditProofReportsTimeoutsDuringAuditProofRequests) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndTimeout("0.123456.999999.tree.ct.test.");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", 123456, 999999, &callback);
+  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_TIMED_OUT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
@@ -584,59 +538,33 @@ TEST_P(LogDnsClientTest, IgnoresLatestDnsConfigIfInvalid) {
   // Let the DNS config change propogate.
   base::RunLoop().RunUntilIdle();
   EXPECT_THAT(dns_client->GetConfig()->nameservers, Not(IsEmpty()));
 }
 
-TEST_P(LogDnsClientTest, CanPerformLeafIndexQueriesInParallel) {
-  // Test that leaf index queries can be performed in parallel.
-  constexpr size_t kNumOfParallelQueries = 3;
-  ASSERT_THAT(kNumOfParallelQueries, AllOf(Le(arraysize(kLeafHashes)),
-                                           Le(arraysize(kBase32LeafHashes))))
-      << "Not enough test data for this many parallel queries";
-
-  std::unique_ptr<LogDnsClient> log_client =
-      CreateLogDnsClient(kNumOfParallelQueries);
-  MockLeafIndexCallback callbacks[kNumOfParallelQueries];
-
-  // Expect multiple queries.
-  for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
-    mock_dns_.ExpectLeafIndexRequestAndResponse(
-        kBase32LeafHashes[i], base::Uint64ToString(kLeafIndices[i]));
-  }
-
-  // Start the queries.
-  for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
-    QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHashes[i],
-                        callbacks[i].AsCallback());
-  }
-
-  // Wait for each query to complete and check its results.
-  for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
-    MockLeafIndexCallback& callback = callbacks[i];
-    callback.WaitUntilRun();
-
-    SCOPED_TRACE(testing::Message() << "callbacks[" << i << "]");
-    ASSERT_TRUE(callback.called());
-    EXPECT_THAT(callback.net_error(), IsOk());
-    EXPECT_THAT(callback.leaf_index(), Eq(kLeafIndices[i]));
-  }
-}
-
-TEST_P(LogDnsClientTest, CanPerformAuditProofQueriesInParallel) {
-  // Check that 3 audit proof queries can be performed in parallel.
+TEST_P(LogDnsClientTest, CanPerformQueriesInParallel) {
+  // Check that 3 queries can be performed in parallel.
   constexpr size_t kNumOfParallelQueries = 3;
   ASSERT_THAT(kNumOfParallelQueries,
-              AllOf(Le(arraysize(kLeafIndices)), Le(arraysize(kTreeSizes))))
+              AllOf(Le(arraysize(kLeafIndexQnames)),
+                    Le(arraysize(kLeafIndices)), Le(arraysize(kTreeSizes))))
       << "Not enough test data for this many parallel queries";
 
   std::unique_ptr<LogDnsClient> log_client =
       CreateLogDnsClient(kNumOfParallelQueries);
   MockAuditProofCallback callbacks[kNumOfParallelQueries];
 
-  // Each query should require one more DNS request than the last.
-  // This helps to test that parallel queries do not intefere with each other,
-  // e.g. one query causing another to end prematurely.
+  // Expect multiple leaf index requests.
+  for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
+    mock_dns_.ExpectLeafIndexRequestAndResponse(
+        kLeafIndexQnames[i], base::Uint64ToString(kLeafIndices[i]));
+  }
+
+  // Make each query require one more audit proof request than the last, by
+  // increasing the number of nodes in the audit proof by
+  // kMaxProofNodesPerDnsResponse for each query. This helps to test that
+  // parallel queries do not intefere with each other, e.g. one query causing
+  // another to end prematurely.
   std::vector<std::string> audit_proofs[kNumOfParallelQueries];
   for (size_t query_i = 0; query_i < kNumOfParallelQueries; ++query_i) {
     const size_t dns_requests_required = query_i + 1;
     audit_proofs[query_i] = GetSampleAuditProof(dns_requests_required *
                                                 kMaxProofNodesPerDnsResponse);
@@ -670,11 +598,11 @@ TEST_P(LogDnsClientTest, CanPerformAuditProofQueriesInParallel) {
     }
   }
 
   // Start the queries.
   for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
-    QueryAuditProofAsync(log_client.get(), "ct.test", kLeafIndices[i],
+    QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[i],
                          kTreeSizes[i], callbacks[i].AsCallback());
   }
 
   // Wait for each query to complete and check its results.
   for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
@@ -690,93 +618,19 @@ TEST_P(LogDnsClientTest, CanPerformAuditProofQueriesInParallel) {
     // EXPECT_THAT(callback.proof()->tree_size, kTreeSizes[i]);
     EXPECT_THAT(callback.proof()->nodes, Eq(audit_proofs[i]));
   }
 }
 
-TEST_P(LogDnsClientTest, CanPerformLeafIndexAndAuditProofQueriesInParallel) {
-  // Check that a leaf index and audit proof query can be performed in parallel.
-  constexpr size_t kNumOfParallelQueries = 2;
-  std::unique_ptr<LogDnsClient> log_client =
-      CreateLogDnsClient(kNumOfParallelQueries);
-  MockLeafIndexCallback leaf_index_callback;
-  MockAuditProofCallback audit_proof_callback;
-  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
-
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "123456");
-
-  // It should require 3 requests to collect the entire audit proof, as there is
-  // only space for 7 nodes per TXT record. One node is 32 bytes long and the
-  // TXT RDATA can have a maximum length of 255 bytes (255 / 32).
-  mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
-                                               audit_proof.begin(),
-                                               audit_proof.begin() + 7);
-  mock_dns_.ExpectAuditProofRequestAndResponse("7.123456.999999.tree.ct.test.",
-                                               audit_proof.begin() + 7,
-                                               audit_proof.begin() + 14);
-  mock_dns_.ExpectAuditProofRequestAndResponse("14.123456.999999.tree.ct.test.",
-                                               audit_proof.begin() + 14,
-                                               audit_proof.end());
-
-  // Start the queries.
-  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHashes[0],
-                      leaf_index_callback.AsCallback());
-  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
-                       audit_proof_callback.AsCallback());
-
-  // Wait for the queries to complete, then check their results.
-  leaf_index_callback.WaitUntilRun();
-  audit_proof_callback.WaitUntilRun();
-
-  ASSERT_TRUE(leaf_index_callback.called());
-  EXPECT_THAT(leaf_index_callback.net_error(), IsOk());
-  EXPECT_THAT(leaf_index_callback.leaf_index(), Eq(123456u));
-
-  ASSERT_TRUE(audit_proof_callback.called());
-  EXPECT_THAT(audit_proof_callback.net_error(), IsOk());
-  ASSERT_THAT(audit_proof_callback.proof(), NotNull());
-  EXPECT_THAT(audit_proof_callback.proof()->leaf_index, Eq(123456u));
-  // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
-  // EXPECT_THAT(audit_proof_callback.proof()->tree_size, Eq(999999));
-  EXPECT_THAT(audit_proof_callback.proof()->nodes, Eq(audit_proof));
-}
-
-TEST_P(LogDnsClientTest, CanBeThrottledToOneLeafIndexQueryAtATime) {
-  // Check that leaf index queries can be rate-limited to one at a time.
-  // The second query, initiated while the first is in progress, should fail.
-  mock_dns_.ExpectLeafIndexRequestAndResponse(kBase32LeafHashes[0], "123456");
-
-  const size_t max_concurrent_queries = 1;
-  std::unique_ptr<LogDnsClient> log_client =
-      CreateLogDnsClient(max_concurrent_queries);
-
-  // Start the queries.
-  MockLeafIndexCallback callback1;
-  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHashes[0],
-                      callback1.AsCallback());
-  MockLeafIndexCallback callback2;
-  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHashes[1],
-                      callback2.AsCallback());
-
-  callback1.WaitUntilRun();
-  callback2.WaitUntilRun();
-
-  // Check that the first query succeeded.
-  ASSERT_TRUE(callback1.called());
-  EXPECT_THAT(callback1.net_error(), IsOk());
-  EXPECT_THAT(callback1.leaf_index(), Eq(123456u));
-
-  // Check that the second query failed.
-  ASSERT_TRUE(callback2.called());
-  EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
-  EXPECT_THAT(callback2.leaf_index(), Eq(0u));
-}
-
-TEST_P(LogDnsClientTest, CanBeThrottledToOneAuditProofQueryAtATime) {
-  // Check that audit proof queries can be rate-limited to one at a time.
+TEST_P(LogDnsClientTest, CanBeThrottledToOneQueryAtATime) {
+  // Check that queries can be rate-limited to one at a time.
   // The second query, initiated while the first is in progress, should fail.
   const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
 
+  // Expect the first query to send leaf index and audit proof requests, but the
+  // second should not due to throttling.
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
+
   // It should require 3 requests to collect the entire audit proof, as there is
   // only space for 7 nodes per TXT record. One node is 32 bytes long and the
   // TXT RDATA can have a maximum length of 255 bytes (255 / 32).
   // Rate limiting should not interfere with these requests.
   mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
@@ -787,20 +641,20 @@ TEST_P(LogDnsClientTest, CanBeThrottledToOneAuditProofQueryAtATime) {
                                                audit_proof.begin() + 14);
   mock_dns_.ExpectAuditProofRequestAndResponse("14.123456.999999.tree.ct.test.",
                                                audit_proof.begin() + 14,
                                                audit_proof.end());
 
-  const size_t max_concurrent_queries = 1;
+  const size_t kMaxConcurrentQueries = 1;
   std::unique_ptr<LogDnsClient> log_client =
-      CreateLogDnsClient(max_concurrent_queries);
+      CreateLogDnsClient(kMaxConcurrentQueries);
 
   // Start the queries.
   MockAuditProofCallback callback1;
-  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+  QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[0], 999999,
                        callback1.AsCallback());
   MockAuditProofCallback callback2;
-  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+  QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[1], 999999,
                        callback2.AsCallback());
 
   callback1.WaitUntilRun();
   callback2.WaitUntilRun();
 
@@ -817,45 +671,10 @@ TEST_P(LogDnsClientTest, CanBeThrottledToOneAuditProofQueryAtATime) {
   ASSERT_TRUE(callback2.called());
   EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
   EXPECT_THAT(callback2.proof(), IsNull());
 }
 
-TEST_P(LogDnsClientTest, ThrottlingAppliesAcrossQueryTypes) {
-  // Check that queries can be rate-limited to one at a time, regardless of the
-  // type of query. The second query, initiated while the first is in progress,
-  // should fail.
-  mock_dns_.ExpectLeafIndexRequestAndResponse(
-      "D4S6DSV2J743QJZEQMH4UYHEYK7KRQ5JIQOCPMFUHZVJNFGHXACA.hash.ct.test.",
-      "123456");
-
-  const size_t max_concurrent_queries = 1;
-  std::unique_ptr<LogDnsClient> log_client =
-      CreateLogDnsClient(max_concurrent_queries);
-
-  // Start the queries.
-  MockLeafIndexCallback leaf_index_callback;
-  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHashes[0],
-                      leaf_index_callback.AsCallback());
-  MockAuditProofCallback audit_proof_callback;
-  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
-                       audit_proof_callback.AsCallback());
-
-  leaf_index_callback.WaitUntilRun();
-  audit_proof_callback.WaitUntilRun();
-
-  // Check that the first query succeeded.
-  ASSERT_TRUE(leaf_index_callback.called());
-  EXPECT_THAT(leaf_index_callback.net_error(), IsOk());
-  EXPECT_THAT(leaf_index_callback.leaf_index(), Eq(123456u));
-
-  // Check that the second query failed.
-  ASSERT_TRUE(audit_proof_callback.called());
-  EXPECT_THAT(audit_proof_callback.net_error(),
-              IsError(net::ERR_TEMPORARILY_THROTTLED));
-  EXPECT_THAT(audit_proof_callback.proof(), IsNull());
-}
-
 INSTANTIATE_TEST_CASE_P(ReadMode,
                         LogDnsClientTest,
                         ::testing::Values(net::IoMode::ASYNC,
                                           net::IoMode::SYNCHRONOUS));