Support custom socket server in test::RunLoop

Puts this to use in stun_server_unittest.

Bug: webrtc:469327588
Change-Id: I66e904cac27c9fafe67c2061650f85996a6a6964
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/460880
Commit-Queue: Tomas Gunnarsson <tommi@webrtc.org>
Auto-Submit: Evan Shrubsole <eshr@webrtc.org>
Commit-Queue: Evan Shrubsole <eshr@webrtc.org>
Reviewed-by: Tomas Gunnarsson <tommi@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#47287}
diff --git a/p2p/test/stun_server_unittest.cc b/p2p/test/stun_server_unittest.cc
index 95f5298..8a63a14 100644
--- a/p2p/test/stun_server_unittest.cc
+++ b/p2p/test/stun_server_unittest.cc
@@ -20,10 +20,10 @@
 #include "rtc_base/byte_buffer.h"
 #include "rtc_base/socket_address.h"
 #include "rtc_base/test_client.h"
-#include "rtc_base/thread.h"
 #include "rtc_base/virtual_socket_server.h"
 #include "test/create_test_environment.h"
 #include "test/gtest.h"
+#include "test/run_loop.h"
 
 namespace webrtc {
 
@@ -57,7 +57,7 @@
  private:
   const Environment env_ = CreateTestEnvironment();
   VirtualSocketServer ss_;
-  AutoSocketServerThread main_thread_{&ss_};
+  test::RunLoop run_loop_{&ss_};
   StunServer server_{AsyncUDPSocket::Create(env_, server_addr, ss_)};
   TestClient client_{AsyncUDPSocket::Create(env_, client_addr, ss_)};
 };
diff --git a/test/BUILD.gn b/test/BUILD.gn
index dbe5920..5b1b5ef 100644
--- a/test/BUILD.gn
+++ b/test/BUILD.gn
@@ -1098,6 +1098,7 @@
     "../rtc_base:threading",
     "../rtc_base:timeutils",
     "../rtc_base:weak_ptr",
+    "//third_party/abseil-cpp/absl/base:nullability",
     "//third_party/abseil-cpp/absl/functional:any_invocable",
   ]
 }
diff --git a/test/run_loop.cc b/test/run_loop.cc
index 82d46bb..a412147 100644
--- a/test/run_loop.cc
+++ b/test/run_loop.cc
@@ -9,6 +9,7 @@
  */
 #include "test/run_loop.h"
 
+#include "absl/base/nullability.h"
 #include "absl/functional/any_invocable.h"
 #include "api/task_queue/pending_task_safety_flag.h"
 #include "api/task_queue/task_queue_base.h"
@@ -21,7 +22,14 @@
 namespace webrtc {
 namespace test {
 
-RunLoop::RunLoop() : weak_factory_(this) {
+RunLoop::RunLoop() : worker_thread_(&socket_server_), weak_factory_(this) {
+  worker_thread_.WrapCurrent();
+}
+
+RunLoop::RunLoop(SocketServer* absl_nullable custom_ss)
+    : socket_server_(custom_ss),
+      worker_thread_(&socket_server_),
+      weak_factory_(this) {
   worker_thread_.WrapCurrent();
 }
 
@@ -69,28 +77,50 @@
 }
 
 RunLoop::FakeSocketServer::FakeSocketServer() = default;
+
+RunLoop::FakeSocketServer::FakeSocketServer(
+    SocketServer* absl_nullable custom_ss)
+    : custom_ss_(custom_ss) {}
+
 RunLoop::FakeSocketServer::~FakeSocketServer() = default;
 
 void RunLoop::FakeSocketServer::FailNextWait() {
   fail_next_wait_ = true;
 }
 
+void RunLoop::FakeSocketServer::SetMessageQueue(Thread* absl_nullable queue) {
+  if (custom_ss_) {
+    custom_ss_->SetMessageQueue(queue);
+  }
+}
+
 bool RunLoop::FakeSocketServer::Wait(TimeDelta max_wait_duration,
                                      bool process_io) {
   if (fail_next_wait_) {
     fail_next_wait_ = false;
     return false;
   }
+  if (custom_ss_) {
+    return custom_ss_->Wait(max_wait_duration, process_io);
+  }
   return true;
 }
 
-void RunLoop::FakeSocketServer::WakeUp() {}
+void RunLoop::FakeSocketServer::WakeUp() {
+  if (custom_ss_) {
+    custom_ss_->WakeUp();
+  }
+}
 
-Socket* RunLoop::FakeSocketServer::CreateSocket(int family, int type) {
+Socket* absl_nullable RunLoop::FakeSocketServer::CreateSocket(int family,
+                                                              int type) {
+  if (custom_ss_) {
+    return custom_ss_->CreateSocket(family, type);
+  }
   return nullptr;
 }
 
-RunLoop::WorkerThread::WorkerThread(SocketServer* ss)
+RunLoop::WorkerThread::WorkerThread(SocketServer* absl_nullable ss)
     : Thread(ss), tq_setter_(this) {}
 
 }  // namespace test
diff --git a/test/run_loop.h b/test/run_loop.h
index 3050ce6..5ae3394 100644
--- a/test/run_loop.h
+++ b/test/run_loop.h
@@ -12,6 +12,7 @@
 
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/functional/any_invocable.h"
 #include "api/task_queue/task_queue_base.h"
 #include "api/units/time_delta.h"
@@ -30,6 +31,7 @@
 class RunLoop {
  public:
   RunLoop();
+  explicit RunLoop(SocketServer* absl_nullable custom_ss);
   ~RunLoop();
 
   // Returns a pointer to the task queue implementation managed by this RunLoop.
@@ -62,30 +64,33 @@
   class FakeSocketServer : public SocketServer {
    public:
     FakeSocketServer();
+    explicit FakeSocketServer(SocketServer* absl_nullable custom_ss);
     ~FakeSocketServer() override;
 
     void FailNextWait();
 
    private:
+    void SetMessageQueue(Thread* absl_nullable queue) override;
     bool Wait(webrtc::TimeDelta max_wait_duration, bool process_io) override;
     void WakeUp() override;
 
-    Socket* CreateSocket(int family, int type) override;
+    Socket* absl_nullable CreateSocket(int family, int type) override;
 
    private:
     bool fail_next_wait_ = false;
+    SocketServer* const absl_nullable custom_ss_ = nullptr;
   };
 
   class WorkerThread : public Thread {
    public:
-    explicit WorkerThread(SocketServer* ss);
+    explicit WorkerThread(SocketServer* absl_nullable ss);
 
    private:
     CurrentTaskQueueSetter tq_setter_;
   };
 
   FakeSocketServer socket_server_;
-  WorkerThread worker_thread_{&socket_server_};
+  WorkerThread worker_thread_;
   WeakPtrFactory<RunLoop> weak_factory_;
 };