services.net: Use the new ModuleNetworkStackClient

Use the new ModuleNetworkStackClient to retrieve instances IpMemoryStore
& IpClient.

Bug: 145825329
Test: Compiles
Change-Id: I065525fdbd64b9509ef0f47bb35954d267c65fd2

3 files changed, 14 insertions, 12 deletions

diff --git a/services/net/java/android/net/IpMemoryStore.java b/services/net/java/android/net/IpMemoryStore.java
index dcefb537d0a0..8df2e0d08e0e 100644
--- a/services/net/java/android/net/IpMemoryStore.java
+++ b/services/net/java/android/net/IpMemoryStore.java
@@ -18,6 +18,7 @@ package android.net;
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.net.networkstack.ModuleNetworkStackClient;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -41,7 +42,7 @@ public class IpMemoryStore extends IpMemoryStoreClient {
         super(context);
         mService = new CompletableFuture<>();
         mTailNode = new AtomicReference<CompletableFuture<IIpMemoryStore>>(mService);
-        getNetworkStackClient().fetchIpMemoryStore(
+        getModuleNetworkStackClient(context).fetchIpMemoryStore(
                 new IIpMemoryStoreCallbacks.Stub() {
                     @Override
                     public void onIpMemoryStoreFetched(@NonNull final IIpMemoryStore memoryStore) {
@@ -85,8 +86,8 @@ public class IpMemoryStore extends IpMemoryStoreClient {
     }
 
     @VisibleForTesting
-    protected NetworkStackClient getNetworkStackClient() {
-        return NetworkStackClient.getInstance();
+    protected ModuleNetworkStackClient getModuleNetworkStackClient(Context context) {
+        return ModuleNetworkStackClient.getInstance(context);
     }
 
     /** Gets an instance of the memory store */
diff --git a/services/net/java/android/net/ip/IpClientUtil.java b/services/net/java/android/net/ip/IpClientUtil.java
index a3618b47171a..b329aeec4853 100644
--- a/services/net/java/android/net/ip/IpClientUtil.java
+++ b/services/net/java/android/net/ip/IpClientUtil.java
@@ -22,7 +22,7 @@ import android.content.Context;
 import android.net.DhcpResultsParcelable;
 import android.net.Layer2PacketParcelable;
 import android.net.LinkProperties;
-import android.net.NetworkStackClient;
+import android.net.networkstack.ModuleNetworkStackClient;
 import android.os.ConditionVariable;
 
 import java.io.FileDescriptor;
@@ -75,11 +75,11 @@ public class IpClientUtil {
      *
      * <p>This is a convenience method to allow clients to use {@link IpClientCallbacks} instead of
      * {@link IIpClientCallbacks}.
-     * @see {@link NetworkStackClient#makeIpClient(String, IIpClientCallbacks)}
+     * @see {@link ModuleNetworkStackClient#makeIpClient(String, IIpClientCallbacks)}
      */
     public static void makeIpClient(Context context, String ifName, IpClientCallbacks callback) {
-        // TODO: migrate clients and remove context argument
-        NetworkStackClient.getInstance().makeIpClient(ifName, new IpClientCallbacksProxy(callback));
+        ModuleNetworkStackClient.getInstance(context)
+                .makeIpClient(ifName, new IpClientCallbacksProxy(callback));
     }
 
     /**
diff --git a/tests/net/java/android/net/IpMemoryStoreTest.java b/tests/net/java/android/net/IpMemoryStoreTest.java
index b81ca36429ff..442ac5605e90 100644
--- a/tests/net/java/android/net/IpMemoryStoreTest.java
+++ b/tests/net/java/android/net/IpMemoryStoreTest.java
@@ -35,6 +35,7 @@ import android.net.ipmemorystore.IOnStatusListener;
 import android.net.ipmemorystore.NetworkAttributes;
 import android.net.ipmemorystore.NetworkAttributesParcelable;
 import android.net.ipmemorystore.Status;
+import android.net.networkstack.ModuleNetworkStackClient;
 import android.os.RemoteException;
 
 import androidx.test.filters.SmallTest;
@@ -67,7 +68,7 @@ public class IpMemoryStoreTest {
     @Mock
     Context mMockContext;
     @Mock
-    NetworkStackClient mNetworkStackClient;
+    ModuleNetworkStackClient mModuleNetworkStackClient;
     @Mock
     IIpMemoryStore mMockService;
     @Mock
@@ -90,14 +91,14 @@ public class IpMemoryStoreTest {
                 ((IIpMemoryStoreCallbacks) invocation.getArgument(0))
                         .onIpMemoryStoreFetched(mMockService);
                 return null;
-            }).when(mNetworkStackClient).fetchIpMemoryStore(any());
+            }).when(mModuleNetworkStackClient).fetchIpMemoryStore(any());
         } else {
-            doNothing().when(mNetworkStackClient).fetchIpMemoryStore(mCbCaptor.capture());
+            doNothing().when(mModuleNetworkStackClient).fetchIpMemoryStore(mCbCaptor.capture());
         }
         mStore = new IpMemoryStore(mMockContext) {
             @Override
-            protected NetworkStackClient getNetworkStackClient() {
-                return mNetworkStackClient;
+            protected ModuleNetworkStackClient getModuleNetworkStackClient(Context ctx) {
+                return mModuleNetworkStackClient;
             }
         };
     }