1 files changed, 10 insertions, 27 deletions

diff --git a/src/plugins/unittest/counter_test.c b/src/plugins/unittest/counter_test.c
index 24b9e1e386e..71f8f93f94d 100644
--- a/src/plugins/unittest/counter_test.c
+++ b/src/plugins/unittest/counter_test.c
@@ -19,7 +19,7 @@
 #include <vppinfra/error.h>
 
 #include <vlib/counter.h>
-#include <vpp/stats/stat_segment.h>
+#include <vlib/stats/stats.h>
 
 enum
 {
@@ -38,31 +38,10 @@ enum
 static uint64_t
 get_stats_epoch ()
 {
-  stat_segment_main_t *sm = &stat_segment_main;
+  vlib_stats_segment_t *sm = vlib_stats_get_segment ();
   return sm->shared_header->epoch;
 }
 
-/*
- * Return the maximum element count of the vector based on its allocated
- * memory.
- */
-static int
-get_vec_mem_size (void *v, uword data_size)
-{
-  stat_segment_main_t *sm = &stat_segment_main;
-
-  if (v == 0)
-    return 0;
-
-  uword aligned_header_bytes = vec_header_bytes (0);
-  void *p = v - aligned_header_bytes;
-  void *oldheap = clib_mem_set_heap (sm->heap);
-  int mem_size = (clib_mem_size (p) - aligned_header_bytes) / data_size;
-  clib_mem_set_heap (oldheap);
-
-  return mem_size;
-}
-
 /* number of times to repeat the counter expand tests */
 #define EXPAND_TEST_ROUNDS 3
 
@@ -90,8 +69,7 @@ test_simple_counter_expand (vlib_main_t *vm)
       // Check how many elements fit into the counter vector without expanding
       // that. The next validate calls should not increase the stats segment
       // epoch.
-      int mem_size = get_vec_mem_size (counter.counters[0],
-				       sizeof ((counter.counters[0])[0]));
+      int mem_size = vec_max_len (counter.counters[0]);
       for (index = 1; index <= mem_size - 1; index++)
 	{
 	  vlib_validate_simple_counter (&counter, index);
@@ -111,6 +89,9 @@ test_simple_counter_expand (vlib_main_t *vm)
       epoch = new_epoch;
     }
 
+  vlib_free_simple_counter (&counter);
+  vlib_validate_simple_counter (&counter, 0);
+
   return 0;
 }
 
@@ -138,8 +119,7 @@ test_combined_counter_expand (vlib_main_t *vm)
       // Check how many elements fit into the counter vector without expanding
       // that. The next validate calls should not increase the stats segment
       // epoch.
-      int mem_size = get_vec_mem_size (counter.counters[0],
-				       sizeof ((counter.counters[0])[0]));
+      int mem_size = vec_max_len (counter.counters[0]);
       for (index = 1; index <= mem_size - 1; index++)
 	{
 	  vlib_validate_combined_counter (&counter, index);
@@ -159,6 +139,9 @@ test_combined_counter_expand (vlib_main_t *vm)
       epoch = new_epoch;
     }
 
+  vlib_free_combined_counter (&counter);
+  vlib_validate_combined_counter (&counter, 0);
+
   return 0;
 }