From 241d12bc484299614c1d1ebec6c4366e125d0c78 Mon Sep 17 00:00:00 2001
From: baxeaz <86568017+baxeaz@users.noreply.github.com>
Date: Fri, 22 Mar 2024 20:45:51 -0700
Subject: [PATCH] feat: make os_user optional in cleanup_session (#232)

Signed-off-by: Brian Axelson <baxelson@amazon.com>
---
 src/deadline/job_attachments/asset_sync.py         | 14 ++++++++++++--
 src/deadline/job_attachments/exceptions.py         |  6 ++++++
 .../deadline_job_attachments/test_asset_sync.py    | 13 +++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/deadline/job_attachments/asset_sync.py b/src/deadline/job_attachments/asset_sync.py
index 5c73af40..7044a8ee 100644
--- a/src/deadline/job_attachments/asset_sync.py
+++ b/src/deadline/job_attachments/asset_sync.py
@@ -39,7 +39,12 @@
     mount_vfs_from_manifests,
 )
 
-from .exceptions import AssetSyncError, VFSExecutableMissingError, JobAttachmentsS3ClientError
+from .exceptions import (
+    AssetSyncError,
+    VFSExecutableMissingError,
+    JobAttachmentsS3ClientError,
+    VFSOSUserNotSetError,
+)
 from .vfs import VFSProcessManager
 from .models import (
     Attachments,
@@ -566,10 +571,15 @@ def sync_outputs(
         return summary_stats
 
     def cleanup_session(
-        self, session_dir: Path, file_system: JobAttachmentsFileSystem, os_user: str
+        self,
+        session_dir: Path,
+        file_system: JobAttachmentsFileSystem,
+        os_user: Optional[str] = None,
     ):
         if file_system == JobAttachmentsFileSystem.COPIED.value:
             return
+        if not os_user:
+            raise VFSOSUserNotSetError("No os user set - can't clean up vfs session")
         try:
             VFSProcessManager.find_vfs()
             # Shutdown all running Deadline VFS processes since session is complete
diff --git a/src/deadline/job_attachments/exceptions.py b/src/deadline/job_attachments/exceptions.py
index 83938b25..3b4c4ade 100644
--- a/src/deadline/job_attachments/exceptions.py
+++ b/src/deadline/job_attachments/exceptions.py
@@ -148,6 +148,12 @@ class VFSFailedToMountError(JobAttachmentsError):
     """
 
 
+class VFSOSUserNotSetError(JobAttachmentsError):
+    """
+    Exception attempting to use the vfs without an os user
+    """
+
+
 class UnsupportedHashingAlgorithmError(JobAttachmentsError):
     """
     Exception for when an unsupported hashing algorithm is provided.
diff --git a/test/unit/deadline_job_attachments/test_asset_sync.py b/test/unit/deadline_job_attachments/test_asset_sync.py
index 648bcd7c..fbe512a7 100644
--- a/test/unit/deadline_job_attachments/test_asset_sync.py
+++ b/test/unit/deadline_job_attachments/test_asset_sync.py
@@ -22,6 +22,7 @@
 from deadline.job_attachments.exceptions import (
     VFSExecutableMissingError,
     JobAttachmentsS3ClientError,
+    VFSOSUserNotSetError,
 )
 from deadline.job_attachments.models import (
     Attachments,
@@ -877,3 +878,15 @@ def test_cleanup_session_vfs_terminate_called(self, tmp_path):
             )
 
             mock_find_vfs.assert_called_once()
+
+    def test_cleanup_session_virtual_witout_os_user_raises(self, tmp_path):
+        self.default_asset_sync.cleanup_session(
+            session_dir=tmp_path,
+            file_system=JobAttachmentsFileSystem.COPIED,
+        )
+
+        with pytest.raises(VFSOSUserNotSetError):
+            self.default_asset_sync.cleanup_session(
+                session_dir=tmp_path,
+                file_system=JobAttachmentsFileSystem.VIRTUAL,
+            )