1 files changed, 32 insertions, 12 deletions

diff --git a/test/core/util/mock_endpoint.c b/test/core/util/mock_endpoint.c
index 13e0e918fb..bf6d85252a 100644
--- a/test/core/util/mock_endpoint.c
+++ b/test/core/util/mock_endpoint.c
@@ -33,24 +33,28 @@
 
 #include "test/core/util/mock_endpoint.h"
 
+#include <inttypes.h>
+
 #include <grpc/support/alloc.h>
 #include <grpc/support/string_util.h>
+#include "src/core/lib/iomgr/sockaddr.h"
 
 typedef struct grpc_mock_endpoint {
   grpc_endpoint base;
   gpr_mu mu;
-  void (*on_write)(gpr_slice slice);
-  gpr_slice_buffer read_buffer;
-  gpr_slice_buffer *on_read_out;
+  void (*on_write)(grpc_slice slice);
+  grpc_slice_buffer read_buffer;
+  grpc_slice_buffer *on_read_out;
   grpc_closure *on_read;
+  grpc_resource_user *resource_user;
 } grpc_mock_endpoint;
 
 static void me_read(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
-                    gpr_slice_buffer *slices, grpc_closure *cb) {
+                    grpc_slice_buffer *slices, grpc_closure *cb) {
   grpc_mock_endpoint *m = (grpc_mock_endpoint *)ep;
   gpr_mu_lock(&m->mu);
   if (m->read_buffer.count > 0) {
-    gpr_slice_buffer_swap(&m->read_buffer, slices);
+    grpc_slice_buffer_swap(&m->read_buffer, slices);
     grpc_exec_ctx_sched(exec_ctx, cb, GRPC_ERROR_NONE, NULL);
   } else {
     m->on_read = cb;
@@ -60,7 +64,7 @@ static void me_read(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
 }
 
 static void me_write(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
-                     gpr_slice_buffer *slices, grpc_closure *cb) {
+                     grpc_slice_buffer *slices, grpc_closure *cb) {
   grpc_mock_endpoint *m = (grpc_mock_endpoint *)ep;
   for (size_t i = 0; i < slices->count; i++) {
     m->on_write(slices->slices[i]);
@@ -83,11 +87,13 @@ static void me_shutdown(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep) {
     m->on_read = NULL;
   }
   gpr_mu_unlock(&m->mu);
+  grpc_resource_user_shutdown(exec_ctx, m->resource_user);
 }
 
 static void me_destroy(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep) {
   grpc_mock_endpoint *m = (grpc_mock_endpoint *)ep;
-  gpr_slice_buffer_destroy(&m->read_buffer);
+  grpc_slice_buffer_destroy(&m->read_buffer);
+  grpc_resource_user_unref(exec_ctx, m->resource_user);
   gpr_free(m);
 }
 
@@ -95,6 +101,13 @@ static char *me_get_peer(grpc_endpoint *ep) {
   return gpr_strdup("fake:mock_endpoint");
 }
 
+static grpc_resource_user *me_get_resource_user(grpc_endpoint *ep) {
+  grpc_mock_endpoint *m = (grpc_mock_endpoint *)ep;
+  return m->resource_user;
+}
+
+static int me_get_fd(grpc_endpoint *ep) { return -1; }
+
 static grpc_workqueue *me_get_workqueue(grpc_endpoint *ep) { return NULL; }
 
 static const grpc_endpoint_vtable vtable = {
@@ -105,13 +118,20 @@ static const grpc_endpoint_vtable vtable = {
     me_add_to_pollset_set,
     me_shutdown,
     me_destroy,
+    me_get_resource_user,
     me_get_peer,
+    me_get_fd,
 };
 
-grpc_endpoint *grpc_mock_endpoint_create(void (*on_write)(gpr_slice slice)) {
+grpc_endpoint *grpc_mock_endpoint_create(void (*on_write)(grpc_slice slice),
+                                         grpc_resource_quota *resource_quota) {
   grpc_mock_endpoint *m = gpr_malloc(sizeof(*m));
   m->base.vtable = &vtable;
-  gpr_slice_buffer_init(&m->read_buffer);
+  char *name;
+  gpr_asprintf(&name, "mock_endpoint_%" PRIxPTR, (intptr_t)m);
+  m->resource_user = grpc_resource_user_create(resource_quota, name);
+  gpr_free(name);
+  grpc_slice_buffer_init(&m->read_buffer);
   gpr_mu_init(&m->mu);
   m->on_write = on_write;
   m->on_read = NULL;
@@ -119,15 +139,15 @@ grpc_endpoint *grpc_mock_endpoint_create(void (*on_write)(gpr_slice slice)) {
 }
 
 void grpc_mock_endpoint_put_read(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
-                                 gpr_slice slice) {
+                                 grpc_slice slice) {
   grpc_mock_endpoint *m = (grpc_mock_endpoint *)ep;
   gpr_mu_lock(&m->mu);
   if (m->on_read != NULL) {
-    gpr_slice_buffer_add(m->on_read_out, slice);
+    grpc_slice_buffer_add(m->on_read_out, slice);
     grpc_exec_ctx_sched(exec_ctx, m->on_read, GRPC_ERROR_NONE, NULL);
     m->on_read = NULL;
   } else {
-    gpr_slice_buffer_add(&m->read_buffer, slice);
+    grpc_slice_buffer_add(&m->read_buffer, slice);
   }
   gpr_mu_unlock(&m->mu);
 }