Changed: use python-dotenv to write the env file (#525)

* Changed: use python-dotenv to write the env file

* Added: test for env overwriting

---------

Co-authored-by: iftwigs <42752431+iftwigs@users.noreply.github.com>

konfuzio_sdk/api.py

@@ -6,9 +6,11 @@
 import time
 from json import JSONDecodeError
 from operator import itemgetter
+from pathlib import Path
 from typing import Dict, List, Optional, Union
 
 import requests
+from dotenv import set_key
 from requests import HTTPError
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
@@ -89,11 +91,13 @@ def init_env(
     :param file_ending: Ending of file.
     """
     token = _get_auth_token(user, password, host)
+    env_file_path = Path(working_directory, file_ending)
+    # Create env file if it does not exist
+    env_file_path.touch(mode=0o600, exist_ok=True)
 
-    with open(os.path.join(working_directory, file_ending), 'w') as f:
-        f.write(f'KONFUZIO_HOST = {host}\n')
-        f.write(f'KONFUZIO_USER = {user}\n')
-        f.write(f'KONFUZIO_TOKEN = {token}\n')
+    set_key(env_file_path, 'KONFUZIO_HOST', host)
+    set_key(env_file_path, 'KONFUZIO_USER', user)
+    set_key(env_file_path, 'KONFUZIO_TOKEN', token)
 
     print('[SUCCESS] SDK initialized!')
 

setup.py

@@ -1,4 +1,5 @@
 """Setup."""
+
 import subprocess
 import sys
 import textwrap
@@ -94,6 +95,7 @@
         'Pillow>=8.4.0',
         'python-dateutil>=2.8.2',
         'python-decouple>=3.3',
+        'python-dotenv>=1.0,<1.1',
         'requests',
         'regex>=2020.6.8',  # re module but better
         'scikit-learn==1.2.2',

tests/test_api.py

@@ -599,25 +599,6 @@ def text(self):
             _get_auth_token('test', 'test')
             assert 'HTTP Status 500' in context.exception
 
-    @patch('requests.post')
-    def test_patched_init_env(self, function):
-        """Test to run CLI."""
-
-        # mock response
-        class _Response:
-            """Mock requests POST response."""
-
-            status_code = 200
-
-            def json(self):
-                """Mock valid return."""
-                return {'token': 'faketoken'}
-
-        function.return_value = _Response()
-        env_file = '.testenv'
-        assert init_env(user='me', password='pw', file_ending=env_file)
-        os.remove(os.path.join(os.getcwd(), env_file))
-
     @patch('konfuzio_sdk.api.konfuzio_session')
     @patch('konfuzio_sdk.api.get_extraction_ais_list_url')
     @patch('konfuzio_sdk.api.get_splitting_ais_list_url')
@@ -715,3 +696,38 @@ def test_init_env():
     """Test to write env file."""
     with pytest.raises(PermissionError, match='Your credentials are not correct'):
         init_env(user='user', password='ABCD', working_directory=BASE_DIR, file_ending='x.env')
+
+
+@patch('requests.post')
+def test_patched_init_env(function):
+    """Test to run CLI."""
+
+    # mock response
+    class _Response:
+        """Mock requests POST response."""
+
+        status_code = 200
+
+        def json(self):
+            """Mock valid return."""
+            return {'token': 'faketoken'}
+
+    function.return_value = _Response()
+    env_file = '.testenv'
+
+    # Create .testenv with some dummy variables
+    with open(os.path.join(os.getcwd(), env_file), 'w') as f:
+        f.write('TESTVAR1=test\n')
+        f.write('KONFUZIO_HOST=testdomain.com\n')
+
+    assert init_env(user='me', password='pw', host='https://testing.konfuzio.com', file_ending=env_file)
+
+    with open(os.path.join(os.getcwd(), env_file), 'r') as f:
+        env_content = f.read()
+
+    # TESTVAR should still be in the env file
+    assert 'TESTVAR1=test' in env_content
+    # KONFUZIO_HOST should be overwritten
+    assert "KONFUZIO_HOST='https://testing.konfuzio.com'" in env_content
+
+    os.remove(os.path.join(os.getcwd(), env_file))