diff --git a/provider/legal-azure/src/test/java/org/opengroup/osdu/legal/azure/jobs/LegalTagPublisherImplTest.java b/provider/legal-azure/src/test/java/org/opengroup/osdu/legal/azure/jobs/LegalTagPublisherImplTest.java
index c48f9df5404bde51eea18081b031aa994af0ed3c..5f71d4104d2c2b1818f805dc84ac0ed87f19f17b 100644
--- a/provider/legal-azure/src/test/java/org/opengroup/osdu/legal/azure/jobs/LegalTagPublisherImplTest.java
+++ b/provider/legal-azure/src/test/java/org/opengroup/osdu/legal/azure/jobs/LegalTagPublisherImplTest.java
@@ -27,6 +27,7 @@ import org.mockito.ArgumentCaptor;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
+import org.opengroup.osdu.azure.servicebus.ITopicClientFactory;
 import org.opengroup.osdu.core.common.logging.JaxRsDpsLog;
 import org.opengroup.osdu.core.common.model.http.DpsHeaders;
 import org.opengroup.osdu.core.common.model.legal.StatusChangedTags;
@@ -42,10 +43,14 @@ public class LegalTagPublisherImplTest {
     private static final String DATA_PARTITION_WITH_FALLBACK_ACCOUNT_ID = "data-partition-account-id";
     private static final String CORRELATION_ID = "correlation-id";
     private static final String USER_EMAIL = "user@email.com";
+    private static final String PARTITION_ID = "partition-id";
 
     @Mock
     private JaxRsDpsLog logger;
 
+    @Mock
+    private ITopicClientFactory topicClientFactory;
+
     @Mock
     private TopicClient topicClient;
 
@@ -57,12 +62,11 @@ public class LegalTagPublisherImplTest {
 
     @Before
     public void init() throws ServiceBusException, InterruptedException {
-        lenient().doNothing().when(topicClient).send(any());
-        lenient().doNothing().when(logger).error(anyString(), (Exception) any());
-        lenient().doNothing().when(logger).info(anyString());
         doReturn(DATA_PARTITION_WITH_FALLBACK_ACCOUNT_ID).when(headers).getPartitionIdWithFallbackToAccountId();
         doReturn(CORRELATION_ID).when(headers).getCorrelationId();
         doReturn(USER_EMAIL).when(headers).getUserEmail();
+        doReturn(PARTITION_ID).when(headers).getPartitionId();
+        doReturn(topicClient).when(topicClientFactory).getClient(eq(PARTITION_ID), any());
     }
 
     @Test