From 848dbfb8a403ffe853e175c9dd1cbb6a6019a84c Mon Sep 17 00:00:00 2001
From: imdhemy <imdhemy@gmail.com>
Date: Sun, 25 Feb 2024 17:35:28 +0100
Subject: [PATCH] chore: add more tests

---
 .../ConnectionPool/SniffingConnectionPool.php |  4 +-
 .../SniffingConnectionPoolTest.php            | 97 +++++++++++++++++++
 2 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/src/OpenSearch/ConnectionPool/SniffingConnectionPool.php b/src/OpenSearch/ConnectionPool/SniffingConnectionPool.php
index 296b0e45..ed7fdd35 100644
--- a/src/OpenSearch/ConnectionPool/SniffingConnectionPool.php
+++ b/src/OpenSearch/ConnectionPool/SniffingConnectionPool.php
@@ -21,7 +21,7 @@
 
 namespace OpenSearch\ConnectionPool;
 
-use Exception;
+use OpenSearch\Common\Exceptions\Curl\OperationTimeoutException;
 use OpenSearch\Common\Exceptions\NoNodesAvailableException;
 use OpenSearch\ConnectionPool\Selectors\SelectorInterface;
 use OpenSearch\Connections\Connection;
@@ -123,7 +123,7 @@ private function sniffConnection(Connection $connection): bool
     {
         try {
             $response = $connection->sniff();
-        } catch (Exception $exception) {
+        } catch (OperationTimeoutException $exception) {
             return false;
         }
 
diff --git a/tests/ConnectionPool/SniffingConnectionPoolTest.php b/tests/ConnectionPool/SniffingConnectionPoolTest.php
index 59a63bab..5935b14e 100644
--- a/tests/ConnectionPool/SniffingConnectionPoolTest.php
+++ b/tests/ConnectionPool/SniffingConnectionPoolTest.php
@@ -15,6 +15,7 @@
 
 namespace OpenSearch\Tests\ConnectionPool;
 
+use OpenSearch\Common\Exceptions\Curl\OperationTimeoutException;
 use OpenSearch\Common\Exceptions\NoNodesAvailableException;
 use OpenSearch\ConnectionPool\Selectors\RoundRobinSelector;
 use OpenSearch\ConnectionPool\SniffingConnectionPool;
@@ -179,6 +180,102 @@ public function itShouldFailIfAllNodesAreDown(): void
         $connectionPool->nextConnection();
     }
 
+    /** @test */
+    public function sniffShouldStopIfAllSniffRequestsFail(): void
+    {
+        $connection = $this->createMock(Connection::class);
+        $connection->method('isAlive')->willReturn(false);
+        $connection->method('sniff')->willThrowException(new OperationTimeoutException());
+        $selector = new RoundRobinSelector();
+        $connectionFactory = $this->createMock(ConnectionFactoryInterface::class);
+
+        $connectionPool = new SniffingConnectionPool(
+            [$connection],
+            $selector,
+            $connectionFactory,
+            ['sniffingInterval' => 0]
+        );
+
+        $this->expectException(NoNodesAvailableException::class);
+        $connectionFactory->expects($this->never())->method('create');
+
+        $connectionPool->nextConnection();
+    }
+
+    /** @test */
+    public function sniffShouldStopIfNodesAreEmpty(): void
+    {
+        $clusterState = $this->clusterState(0);
+        $connection = $this->createMock(Connection::class);
+        $connection->method('isAlive')->willReturn(false);
+        $connection->method('sniff')->willReturn($clusterState);
+        $selector = new RoundRobinSelector();
+        $connectionFactory = $this->createMock(ConnectionFactoryInterface::class);
+
+        $connectionPool = new SniffingConnectionPool(
+            [$connection],
+            $selector,
+            $connectionFactory,
+            ['sniffingInterval' => 0]
+        );
+
+        $this->expectException(NoNodesAvailableException::class);
+        $connectionFactory->expects($this->never())->method('create');
+
+        $connectionPool->nextConnection();
+    }
+
+    /** @test */
+    public function itShouldNotSniffBeforeScheduledSniffTime(): void
+    {
+        $connection = $this->createMock(Connection::class);
+        $connection->method('isAlive')->willReturn(false);
+        $connection->method('sniff')->willReturn($this->clusterState(2));
+        $selector = new RoundRobinSelector();
+        $connectionFactory = $this->createMock(ConnectionFactoryInterface::class);
+
+        $connectionPool = new SniffingConnectionPool(
+            [$connection],
+            $selector,
+            $connectionFactory,
+            ['sniffingInterval' => 300]
+        );
+
+        $connectionFactory->expects($this->never())->method('create');
+        $this->expectException(NoNodesAvailableException::class);
+
+        $connectionPool->nextConnection();
+    }
+
+    /** @test */
+    public function scheduleCheck(): void
+    {
+        $clusterState = $this->clusterState(2);
+        $firstConnection = $this->createMock(Connection::class);
+        $firstConnection->method('isAlive')->willReturn(true);
+        $firstConnection->method('sniff')->willReturn($clusterState);
+        $secondConnection = $this->createMock(Connection::class);
+        $secondConnection->method('isAlive')->willReturn(true);
+        $selector = $this->createMock(RoundRobinSelector::class);
+        $selector->expects($this->exactly(2))->method('select')->willReturnOnConsecutiveCalls(
+            $firstConnection,
+            $secondConnection
+        );
+        $connectionFactory = $this->createMock(ConnectionFactoryInterface::class);
+        $connectionFactory->method('create')->willReturnOnConsecutiveCalls($firstConnection, $secondConnection);
+
+        $connectionPool = new SniffingConnectionPool(
+            [$firstConnection],
+            $selector,
+            $connectionFactory,
+            ['sniffingInterval' => 300]
+        );
+
+        $connectionPool->scheduleCheck();
+
+        $this->assertSame($secondConnection, $connectionPool->nextConnection());
+    }
+
     private function clusterState(int $numberOfNodes): array
     {
         $clusterState = ['nodes' => []];