addressing code comments

Author: AGhafaryy

nodestream/pipeline/extractors/credential_utils.py

@@ -7,26 +7,19 @@
 import pytz
 from botocore.credentials import RefreshableCredentials
 from botocore.session import Session
-from nodestream.file_io import LazyLoadedArgument
 
 
 class AwsClientFactory:
     def __init__(
         self,
-        assume_role_arn: Optional[Union[LazyLoadedArgument,str]] = None,
-        assume_role_external_id: Optional[Union[LazyLoadedArgument,str]] = None,
+        assume_role_arn: Optional[str] = None,
+        assume_role_external_id: Optional[str] = None,
         session_ttl: int = 3000,
         **boto_session_args
     ) -> None:
-        if isinstance(assume_role_arn, LazyLoadedArgument):
-            self.assume_role_arn = assume_role_arn.get_value()
-        else:
-            self.assume_role_arn = assume_role_arn
-        if isinstance(assume_role_external_id, LazyLoadedArgument):
-            self.assume_role_external_id = assume_role_external_id.get_value()
-        else:
-            self.assume_role_external_id = assume_role_external_id
 
+        self.assume_role_arn = assume_role_arn
+        self.assume_role_external_id = assume_role_external_id
         self.session_args = self._init_session_args(**boto_session_args)
         self.session_ttl = session_ttl
 

nodestream/project/storage.py

@@ -12,13 +12,11 @@ class StoreConfiguration:
     hmac_key: Optional[Union[LazyLoadedArgument,str]] = None
 
     def initialize(self) -> ObjectStore:
-        store = ObjectStore.from_file_arguments(self.storage_type, **self.arguments)
+        args = LazyLoadedArgument.resolve_if_needed(self.arguments)
+        store = ObjectStore.from_file_arguments(self.storage_type, **args)
         if self.hmac_key:
-            if isinstance(self.hmac_key, LazyLoadedArgument):
-                resolved_key = self.hmac_key.get_value()
-                return store.signed(Signer.hmac(resolved_key))
-            else:
-                return store.signed(Signer.hmac(self.hmac_key))
+            resolved_key = LazyLoadedArgument.resolve_if_needed(self.hmac_key)
+            return store.signed(Signer.hmac(resolved_key))
         return store
 
     def to_file_data(self):
@@ -40,10 +38,10 @@ def from_file_data(data):
 
     @staticmethod
     def describe_yaml_schema():
-        from schema import Optional, Schema
+        from schema import Optional, Schema, Or
 
         return Schema(
-            {"name": str, "type": str, Optional("hmac_key"): LazyLoadedArgument, Optional(str): object}
+            {"name": str, "type": str, Optional("hmac_key"): Or(LazyLoadedArgument, str, only_one=True), Optional(str): object}
         )
 
 

tests/unit/pipeline/extractors/stores/aws/test_credential_utils.py

@@ -6,8 +6,6 @@
 from freezegun import freeze_time
 from hamcrest import assert_that, equal_to, has_key, not_
 
-from nodestream.file_io import LazyLoadedArgument
-from unittest.mock import Mock
 
 
 @pytest.fixture
@@ -148,37 +146,3 @@ def test_make_client(mocker, client_without_role):
         return_value=session
     )
     client_without_role.make_client("s3")
-
-def create_mock_lazy_loaded_argument(value):
-    mock = Mock(spec=LazyLoadedArgument)
-    mock.get_value.return_value = value
-    return mock
-
-def test_init_with_lazy_loaded_role_arn():
-    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
-
-    lazy_role_arn = create_mock_lazy_loaded_argument("arn:aws:iam::123456789012:role/test")
-    client = AwsClientFactory(assume_role_arn=lazy_role_arn)
-
-    assert_that(client.assume_role_arn, equal_to("arn:aws:iam::123456789012:role/test"))
-
-def test_init_with_lazy_loaded_external_id():
-    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
-
-    lazy_external_id = create_mock_lazy_loaded_argument("test-external-id")
-    client = AwsClientFactory(assume_role_external_id=lazy_external_id)
-
-    assert_that(client.assume_role_external_id, equal_to("test-external-id"))
-
-def test_init_with_both_lazy_loaded_arguments():
-    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
-
-    lazy_role_arn = create_mock_lazy_loaded_argument("arn:aws:iam::123456789012:role/test")
-    lazy_external_id = create_mock_lazy_loaded_argument("test-external-id")
-    client = AwsClientFactory(
-        assume_role_arn=lazy_role_arn,
-        assume_role_external_id=lazy_external_id
-    )
-
-    assert_that(client.assume_role_arn, equal_to("arn:aws:iam::123456789012:role/test"))
-    assert_that(client.assume_role_external_id, equal_to("test-external-id"))

tests/unit/project/test_storage.py

@@ -96,15 +96,18 @@ def test_storage_configuration_to_file_data(store_config):
     assert_that(file_data, equal_to(expected_data))
 
 
-def test_store_configuration_initialize_with_lazy_hmac(tmp_path):
+def test_store_configuration_initialize_with_lazy_hmac():
     expected_hmac = "dvHdCrVbRPp1HcmWX78Ryw=="
     mock_lazy_hmac = Mock(spec=LazyLoadedArgument)
     mock_lazy_hmac.get_value.return_value = expected_hmac
+    expected_arg = "testarglol"
+    mock_lazy_arg = Mock(spec=LazyLoadedArgument)
+    mock_lazy_arg.get_value.return_value = expected_arg
 
     store_config = StoreConfiguration(
         name="test-store",
         storage_type="local",
-        arguments={"root": tmp_path},
+        arguments={"root": mock_lazy_arg},
         hmac_key=mock_lazy_hmac
     )
 
@@ -113,4 +116,5 @@ def test_store_configuration_initialize_with_lazy_hmac(tmp_path):
         base64.b64encode(store.signer.key).decode(),
         equal_to(expected_hmac)
     )
+    assert_that(store.store.root, equal_to(expected_arg))
     mock_lazy_hmac.get_value.assert_called_once()