From 0df45d6da53baee07369aa38ec9ce46990c589c9 Mon Sep 17 00:00:00 2001
From: Michael Churilov <michael.churilov@gmail.com>
Date: Thu, 25 Apr 2024 21:35:38 -0400
Subject: [PATCH] Adding amqps as a valid broker scheme

Issue here: https://github.com/mher/flower/issues/1369

Flower does not take into account the `BROKER_API` config setting right
now if the scheme is set to `amqps`. It errors.

This seems like it should not be the case, as there is no change
required for the RabbitMQ API to work, regardless of amqp or amqps.
---
 flower/api/tasks.py             |  2 +-
 flower/utils/broker.py          |  2 ++
 flower/views/broker.py          |  2 +-
 tests/unit/utils/test_broker.py | 46 ++++++++++++++++++++++++++-------
 4 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/flower/api/tasks.py b/flower/api/tasks.py
index 730c290e..f7d28fbc 100644
--- a/flower/api/tasks.py
+++ b/flower/api/tasks.py
@@ -393,7 +393,7 @@ async def get(self):
         app = self.application
 
         http_api = None
-        if app.transport == 'amqp' and app.options.broker_api:
+        if (app.transport == 'amqp' or app.transport == 'amqps') and app.options.broker_api:
             http_api = app.options.broker_api
 
         broker = Broker(app.capp.connection().as_uri(include_password=True),
diff --git a/flower/utils/broker.py b/flower/utils/broker.py
index 51255c4a..8aded026 100644
--- a/flower/utils/broker.py
+++ b/flower/utils/broker.py
@@ -230,6 +230,8 @@ def __new__(cls, broker_url, *args, **kwargs):
         scheme = urlparse(broker_url).scheme
         if scheme == 'amqp':
             return RabbitMQ(broker_url, *args, **kwargs)
+        if scheme == 'amqps':
+            return RabbitMQ(broker_url, *args, **kwargs)
         if scheme == 'redis':
             return Redis(broker_url, *args, **kwargs)
         if scheme == 'rediss':
diff --git a/flower/views/broker.py b/flower/views/broker.py
index 75f6c9b3..043bb937 100644
--- a/flower/views/broker.py
+++ b/flower/views/broker.py
@@ -14,7 +14,7 @@ async def get(self):
         app = self.application
 
         http_api = None
-        if app.transport == 'amqp' and app.options.broker_api:
+        if (app.transport == 'amqp' or app.transport == 'amqps') and app.options.broker_api:
             http_api = app.options.broker_api
 
         try:
diff --git a/tests/unit/utils/test_broker.py b/tests/unit/utils/test_broker.py
index f3d3ee1b..43370e62 100644
--- a/tests/unit/utils/test_broker.py
+++ b/tests/unit/utils/test_broker.py
@@ -9,7 +9,7 @@
 broker.redis = MagicMock()
 
 
-class TestRabbitMQ(unittest.TestCase):
+class TestRabbitMQwithAMQP(unittest.TestCase):
     def test_init(self):
         b = Broker('amqp://', '')
         self.assertTrue(isinstance(b, RabbitMQ))
@@ -40,18 +40,46 @@ def test_url_defaults_rabbitmq(self):
             self.assertEqual('guest', b.username)
             self.assertEqual('guest', b.password)
 
-    def test_url_defaults_redis(self):
-        for url in ['redis://', 'redis://localhost', 'redis://localhost/0']:
-            b = Redis(url, '')
+    def test_invalid_http_api(self):
+        with self.assertLogs('', level='ERROR') as cm:
+            RabbitMQ('amqp://user:pass@host:10000/vhost', http_api='ftp://')
+            self.assertEqual(['ERROR:flower.utils.broker:Invalid broker api url: ftp://'], cm.output)
+
+
+class TestRabbitMQwithAMQPS(unittest.TestCase):
+    def test_init(self):
+        b = Broker('amqps://', '')
+        self.assertTrue(isinstance(b, RabbitMQ))
+        self.assertFalse(isinstance(b, Redis))
+
+    def test_url(self):
+        b = RabbitMQ('amqps://user:pass@host:10000/vhost', '')
+        self.assertEqual('host', b.host)
+        self.assertEqual(10000, b.port)
+        self.assertEqual('vhost', b.vhost)
+        self.assertEqual('user', b.username)
+        self.assertEqual('pass', b.password)
+
+    def test_url_vhost_slash(self):
+        b = RabbitMQ('amqps://user:pass@host:10000//', '')
+        self.assertEqual('host', b.host)
+        self.assertEqual(10000, b.port)
+        self.assertEqual('/', b.vhost)
+        self.assertEqual('user', b.username)
+        self.assertEqual('pass', b.password)
+
+    def test_url_defaults_rabbitmq(self):
+        for url in ['amqps://', 'amqp://localhost']:
+            b = RabbitMQ(url, '')
             self.assertEqual('localhost', b.host)
-            self.assertEqual(6379, b.port)
-            self.assertEqual(0, b.vhost)
-            self.assertEqual(None, b.username)
-            self.assertEqual(None, b.password)
+            self.assertEqual(15672, b.port)
+            self.assertEqual('/', b.vhost)
+            self.assertEqual('guest', b.username)
+            self.assertEqual('guest', b.password)
 
     def test_invalid_http_api(self):
         with self.assertLogs('', level='ERROR') as cm:
-            RabbitMQ('amqp://user:pass@host:10000/vhost', http_api='ftp://')
+            RabbitMQ('amqps://user:pass@host:10000/vhost', http_api='ftp://')
             self.assertEqual(['ERROR:flower.utils.broker:Invalid broker api url: ftp://'], cm.output)