Skip to content

Commit

Permalink
Fix ML tests
Browse files Browse the repository at this point in the history
Signed-off-by: Kim Pepper <[email protected]>
  • Loading branch information
kimpepper committed Jan 8, 2025
1 parent 5855571 commit 4dd418d
Show file tree
Hide file tree
Showing 4 changed files with 114 additions and 134 deletions.
Original file line number Diff line number Diff line change
@@ -1,29 +1,28 @@
<?php

declare(strict_types=1);

namespace OpenSearch\Tests\Namespaces;

use OpenSearch\EndpointFactory;
use OpenSearch\EndpointFactoryInterface;
use OpenSearch\Endpoints\Ml\CreateConnector;
use OpenSearch\Endpoints\Ml\DeleteConnector;
use OpenSearch\Endpoints\Ml\GetConnector;
use OpenSearch\Endpoints\Ml\GetConnectors;
use OpenSearch\Endpoints\Ml\DeleteModelGroup;
use OpenSearch\Endpoints\Ml\GetModelGroups;
use OpenSearch\Endpoints\Ml\RegisterModelGroup;
use OpenSearch\Endpoints\Ml\UpdateModelGroup;
use OpenSearch\Endpoints\Ml\DeleteModel;
use OpenSearch\Endpoints\Ml\DeleteModelGroup;
use OpenSearch\Endpoints\Ml\DeployModel;
use OpenSearch\Endpoints\Ml\GetConnector;
use OpenSearch\Endpoints\Ml\GetConnectors;
use OpenSearch\Endpoints\Ml\GetModel;
use OpenSearch\Endpoints\Ml\SearchModels;
use OpenSearch\Endpoints\Ml\GetModelGroups;
use OpenSearch\Endpoints\Ml\GetTask;
use OpenSearch\Endpoints\Ml\Predict;
use OpenSearch\Endpoints\Ml\RegisterModel;
use OpenSearch\Endpoints\Ml\RegisterModelGroup;
use OpenSearch\Endpoints\Ml\SearchModels;
use OpenSearch\Endpoints\Ml\UndeployModel;
use OpenSearch\Endpoints\Ml\GetTask;
use OpenSearch\Endpoints\Sql\Query;
use OpenSearch\Endpoints\Ml\UpdateModelGroup;
use OpenSearch\Namespaces\MlNamespace;
use OpenSearch\Namespaces\SqlNamespace;
use OpenSearch\Transport;
use OpenSearch\TransportInterface;
use PHPUnit\Framework\TestCase;

/**
Expand All @@ -43,15 +42,13 @@ public function testCreatingConnector(): void
$endpointFactory->method('getEndpoint')
->willReturn(new CreateConnector());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/connectors/_create', [], [
'foo' => 'bar',
]);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/connectors/_create', [], [
'foo' => 'bar',
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->createConnector([
'body' => [
Expand All @@ -67,16 +64,15 @@ public function testGetConnector(): void
$endpointFactory->method('getEndpoint')
->willReturn(new GetConnector());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('GET', '/_plugins/_ml/connectors/foobar', [], null);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('GET', '/_plugins/_ml/connectors/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->getConnector([
'id' => 'foobar'
'id' => 'foobar',
'connector_id' => 'foobar'
]);
}

Expand All @@ -87,18 +83,16 @@ public function testGetConnectors(): void
$endpointFactory->method('getEndpoint')
->willReturn(new GetConnectors());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/connectors/_search', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
]);

$transport->method('resultOrFuture')
->willReturn([]);
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->getConnectors([
'body' => [
Expand All @@ -117,13 +111,11 @@ public function testDeleteConnector(): void
$endpointFactory->method('getEndpoint')
->willReturn(new DeleteConnector());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('DELETE', '/_plugins/_ml/connectors/foobar', [], null);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('DELETE', '/_plugins/_ml/connectors/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->deleteConnector([
'connector_id' => 'foobar'
Expand All @@ -137,15 +129,13 @@ public function testRegisterModelGroup(): void
$endpointFactory->method('getEndpoint')
->willReturn(new RegisterModelGroup());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/model_groups/_register', [], [
'foo' => 'bar',
]);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/model_groups/_register', [], [
'foo' => 'bar',
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->registerModelGroup([
'body' => [
Expand All @@ -161,18 +151,16 @@ public function testGetModelGroups(): void
$endpointFactory->method('getEndpoint')
->willReturn(new GetModelGroups());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/model_groups/_search', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
]);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/model_groups/_search', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->getModelGroups([
'body' => [
Expand All @@ -191,21 +179,20 @@ public function testUpdateModelGroup(): void
$endpointFactory->method('getEndpoint')
->willReturn(new UpdateModelGroup());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('PUT', '/_plugins/_ml/model_groups/foobar', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
]);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('PUT', '/_plugins/_ml/model_groups/foobar', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->updateModelGroup([
'id' => 'foobar',
'model_group_id' => 'foobar',
'body' => [
'query' => [
'match_all' => new \StdClass(),
Expand All @@ -222,13 +209,11 @@ public function testDeleteModelGroup(): void
$endpointFactory->method('getEndpoint')
->willReturn(new DeleteModelGroup());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('DELETE', '/_plugins/_ml/model_groups/foobar', [], null);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('DELETE', '/_plugins/_ml/model_groups/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->deleteModelGroup([
'id' => 'foobar'
Expand All @@ -242,15 +227,13 @@ public function testRegisterModel(): void
$endpointFactory->method('getEndpoint')
->willReturn(new RegisterModel());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/models/_register', [], [
'foo' => 'bar',
]);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/models/_register', [], [
'foo' => 'bar',
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->registerModel([
'body' => [
Expand All @@ -265,13 +248,11 @@ public function testGetModel(): void
$endpointFactory->method('getEndpoint')
->willReturn(new GetModel());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('GET', '/_plugins/_ml/models/foobar', [], null);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('GET', '/_plugins/_ml/models/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->getModel([
'id' => 'foobar',
Expand All @@ -284,18 +265,16 @@ public function testSearchModels(): void
$endpointFactory->method('getEndpoint')
->willReturn(new SearchModels());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
$transport->method('sendRequest')
->with('GET', '/_plugins/_ml/models/_search', [], [
'query' => [
'match_all' => new \StdClass(),
],
'size' => 1000,
]);

$transport->method('resultOrFuture')
->willReturn([]);
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->searchModels([
'body' => [
Expand All @@ -313,13 +292,11 @@ public function testDeployModel(): void
$endpointFactory->method('getEndpoint')
->willReturn(new DeployModel());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/models/foobar/_deploy', [], null);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/models/foobar/_deploy', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->deployModel([
'model_id' => 'foobar',
Expand All @@ -332,13 +309,11 @@ public function testUnDeployModel(): void
$endpointFactory->method('getEndpoint')
->willReturn(new UndeployModel());

$transport = $this->createMock(Transport::class);
$transport = $this->createMock(TransportInterface::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/models/foobar/_undeploy', [], null);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/models/foobar/_undeploy', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->undeployModel([
'model_id' => 'foobar',
Expand All @@ -351,13 +326,11 @@ public function testDeleteModel(): void
$endpointFactory->method('getEndpoint')
->willReturn(new DeleteModel());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('DELETE', '/_plugins/_ml/models/foobar', [], null);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('DELETE', '/_plugins/_ml/models/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->deleteModel([
'id' => 'foobar',
Expand All @@ -370,21 +343,21 @@ public function testPredict(): void
$endpointFactory->method('getEndpoint')
->willReturn(new Predict());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('POST', '/_plugins/_ml/models/foobar/_predict', [], [
'foo' => 'bar',
]);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('POST', '/_plugins/_ml/_predict/algo/model', [], [
'foo' => 'bar',
])
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->predict([
'id' => 'foobar',
'body' => [
'foo' => 'bar',
]
'id' => 'foobar',
'body' => [
'foo' => 'bar',
],
'algorithm_name' => 'algo',
'model_id' => 'model',
]);
}

Expand All @@ -394,13 +367,11 @@ public function testGetTask(): void
$endpointFactory->method('getEndpoint')
->willReturn(new GetTask());

$transport = $this->createMock(Transport::class);

$transport->method('performRequest')
->with('GET', '/_plugins/_ml/tasks/foobar', [], null);
$transport = $this->createMock(TransportInterface::class);

$transport->method('resultOrFuture')
->willReturn([]);
$transport->method('sendRequest')
->with('GET', '/_plugins/_ml/tasks/foobar', [], null)
->willReturn([]);

(new MlNamespace($transport, $endpointFactory))->getTask([
'id' => 'foobar',
Expand Down
Loading

0 comments on commit 4dd418d

Please sign in to comment.