From c1902bb558132cfa8cfdeea7265fde454ff55042 Mon Sep 17 00:00:00 2001
From: Matthew Kotila <matthew.r.kotila@gmail.com>
Date: Thu, 10 Aug 2023 22:31:22 +0000
Subject: [PATCH] Address feedback

---
 src/c++/perf_analyzer/inference_profiler.cc   |  3 ++
 .../perf_analyzer/test_inference_profiler.cc  | 51 ++++++++++++++-----
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/c++/perf_analyzer/inference_profiler.cc b/src/c++/perf_analyzer/inference_profiler.cc
index cb870f6f7..da5acf255 100644
--- a/src/c++/perf_analyzer/inference_profiler.cc
+++ b/src/c++/perf_analyzer/inference_profiler.cc
@@ -1298,6 +1298,9 @@ InferenceProfiler::ValidLatencyMeasurement(
           (request_end_ns <= valid_range.second)) {
         valid_latencies->push_back(request_end_ns - request_start_ns);
         response_count += request_record.response_times_.size();
+        if (request_record.has_null_last_response_) {
+          response_count--;
+        }
         erase_indices.push_back(i);
         if (request_record.sequence_end_) {
           valid_sequence_count++;
diff --git a/src/c++/perf_analyzer/test_inference_profiler.cc b/src/c++/perf_analyzer/test_inference_profiler.cc
index 2c4a47d20..5a41319bd 100644
--- a/src/c++/perf_analyzer/test_inference_profiler.cc
+++ b/src/c++/perf_analyzer/test_inference_profiler.cc
@@ -885,14 +885,41 @@ TEST_CASE(
         0, false, 0, false)};
 
     auto request2_timestamp{clock_epoch + std::chrono::nanoseconds(4)};
-    auto response3_timestamp{clock_epoch + std::chrono::nanoseconds(5)};
-    auto response4_timestamp{clock_epoch + std::chrono::nanoseconds(6)};
-    auto response5_timestamp{clock_epoch + std::chrono::nanoseconds(7)};
-    auto request_record2{RequestRecord(
-        request2_timestamp,
-        std::vector<std::chrono::time_point<std::chrono::system_clock>>{
-            response3_timestamp, response4_timestamp, response5_timestamp},
-        0, false, 0, false)};
+    RequestRecord request_record2{};
+    size_t expected_response_count{0};
+
+    SUBCASE("second request has three data responses")
+    {
+      auto response3_timestamp{clock_epoch + std::chrono::nanoseconds(5)};
+      auto response4_timestamp{clock_epoch + std::chrono::nanoseconds(6)};
+      auto response5_timestamp{clock_epoch + std::chrono::nanoseconds(7)};
+      request_record2 = RequestRecord(
+          request2_timestamp,
+          std::vector<std::chrono::time_point<std::chrono::system_clock>>{
+              response3_timestamp, response4_timestamp, response5_timestamp},
+          0, false, 0, false);
+      expected_response_count = 5;
+    }
+    SUBCASE("second request has two data responses and one null response")
+    {
+      auto response3_timestamp{clock_epoch + std::chrono::nanoseconds(5)};
+      auto response4_timestamp{clock_epoch + std::chrono::nanoseconds(6)};
+      auto response5_timestamp{clock_epoch + std::chrono::nanoseconds(7)};
+      request_record2 = RequestRecord(
+          request2_timestamp,
+          std::vector<std::chrono::time_point<std::chrono::system_clock>>{
+              response3_timestamp, response4_timestamp, response5_timestamp},
+          0, false, 0, true);
+      expected_response_count = 4;
+    }
+    SUBCASE("second request has one null response")
+    {
+      request_record2 = RequestRecord(
+          request2_timestamp,
+          std::vector<std::chrono::time_point<std::chrono::system_clock>>{}, 0,
+          false, 0, true);
+      expected_response_count = 2;
+    }
 
     mock_inference_profiler.all_request_records_ = {
         request_record1, request_record2};
@@ -909,7 +936,7 @@ TEST_CASE(
         valid_range, valid_sequence_count, delayed_request_count,
         &valid_latencies, response_count, valid_requests);
 
-    CHECK(response_count == 5);
+    CHECK(response_count == expected_response_count);
   }
   SUBCASE("testing logic relevant to valid request output")
   {
@@ -921,7 +948,7 @@ TEST_CASE(
         request1_timestamp,
         std::vector<std::chrono::time_point<std::chrono::system_clock>>{
             response1_timestamp},
-        0, false, 0)};
+        0, false, 0, false)};
 
     auto request2_timestamp{clock_epoch + std::chrono::nanoseconds(3)};
     auto response2_timestamp{clock_epoch + std::chrono::nanoseconds(4)};
@@ -929,7 +956,7 @@ TEST_CASE(
         request2_timestamp,
         std::vector<std::chrono::time_point<std::chrono::system_clock>>{
             response2_timestamp},
-        0, false, 0)};
+        0, false, 0, false)};
 
     auto request3_timestamp{clock_epoch + std::chrono::nanoseconds(5)};
     auto response3_timestamp{clock_epoch + std::chrono::nanoseconds(6)};
@@ -937,7 +964,7 @@ TEST_CASE(
         request3_timestamp,
         std::vector<std::chrono::time_point<std::chrono::system_clock>>{
             response3_timestamp},
-        0, false, 0)};
+        0, false, 0, false)};
 
     mock_inference_profiler.all_request_records_ = {
         request_record1, request_record2, request_record3};