Fixing tests

Author: jonathan-buttner

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceActionTestCase.java

@@ -57,11 +57,15 @@ public abstract class BaseTransportInferenceActionTestCase<Request extends BaseI
     private BaseTransportInferenceAction<Request> action;
 
     protected static final String serviceId = "serviceId";
-    protected static final TaskType taskType = TaskType.COMPLETION;
+    protected final TaskType taskType;
     protected static final String inferenceId = "inferenceEntityId";
     protected InferenceServiceRegistry serviceRegistry;
     protected InferenceStats inferenceStats;
 
+    public BaseTransportInferenceActionTestCase(TaskType taskType) {
+        this.taskType = taskType;
+    }
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
@@ -377,7 +381,7 @@ protected void mockModelAndServiceRegistry(InferenceService service) {
         when(serviceRegistry.getService(any())).thenReturn(Optional.of(service));
     }
 
-    protected void mockValidLicenseState(){
+    protected void mockValidLicenseState() {
         when(licenseState.isAllowed(InferencePlugin.INFERENCE_API_FEATURE)).thenReturn(true);
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceActionTests.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.inference.InferenceServiceRegistry;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.license.MockLicenseState;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
@@ -20,6 +21,10 @@
 
 public class TransportInferenceActionTests extends BaseTransportInferenceActionTestCase<InferenceAction.Request> {
 
+    public TransportInferenceActionTests() {
+        super(TaskType.COMPLETION);
+    }
+
     @Override
     protected BaseTransportInferenceAction<InferenceAction.Request> createAction(
         TransportService transportService,

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionActionTests.java

@@ -33,6 +33,10 @@
 
 public class TransportUnifiedCompletionActionTests extends BaseTransportInferenceActionTestCase<UnifiedCompletionAction.Request> {
 
+    public TransportUnifiedCompletionActionTests() {
+        super(TaskType.CHAT_COMPLETION);
+    }
+
     @Override
     protected BaseTransportInferenceAction<UnifiedCompletionAction.Request> createAction(
         TransportService transportService,