From 9606e3e269e3b08c72a1acb3ae9f6635ab05bc99 Mon Sep 17 00:00:00 2001
From: Pablo Martin <pablo.martin@lolamarket.com>
Date: Mon, 16 Jan 2023 14:03:12 +0100
Subject: [PATCH] Fetch prefect host and test

---
 lolafect/lolaconfig.py   | 11 +++++++++++
 tests/test_lolaconfig.py | 17 +++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/lolafect/lolaconfig.py b/lolafect/lolaconfig.py
index 515d824..e615d66 100644
--- a/lolafect/lolaconfig.py
+++ b/lolafect/lolaconfig.py
@@ -155,6 +155,17 @@ class LolaConfig:
             "port": self.ENV_DATA["datadw_port"],
         }
 
+    @_needs_env_data
+    def fetch_prefect_host(self, s3_reader=None) -> None:
+        """
+        Read the env file from S3 and store the prefect_host.
+
+        :param s3_reader: a client to fetch files from S3.
+        :return: None
+        """
+
+        self.PREFECT_HOST = self.ENV_DATA["prefect_host"]
+
     def fetch_env_data(self, s3_reader=None) -> None:
         """
         Read the env file from S3 with the default or a passed s3_reader and
diff --git a/tests/test_lolaconfig.py b/tests/test_lolaconfig.py
index 212bc69..64b6094 100644
--- a/tests/test_lolaconfig.py
+++ b/tests/test_lolaconfig.py
@@ -91,3 +91,20 @@ def test_lolaconfig_fetches_dw_creds_properly():
     lolaconfig.fetch_dw_credentials(s3_reader=fake_s3_reader)
 
     assert type(lolaconfig.DW_CREDENTIALS) is dict
+
+
+def test_lolaconfig_fetches_prefect_host_properly():
+    lolaconfig = LolaConfig(flow_name="some-flow")
+
+    fake_s3_reader = SimpleNamespace()
+
+    def mock_read_json_from_s3_file(bucket, key):
+        return {
+            "prefect_host": "some_host",
+        }
+
+    fake_s3_reader.read_json_from_s3_file = mock_read_json_from_s3_file
+
+    lolaconfig.fetch_prefect_host(s3_reader=fake_s3_reader)
+
+    assert type(lolaconfig.PREFECT_HOST) is str