From 9819797434ab6598ca5a5b211734f18ce8886388 Mon Sep 17 00:00:00 2001
From: George Mattimoe <george.mattimoe@octoenergy.com>
Date: Tue, 2 May 2023 16:51:38 +0100
Subject: [PATCH] Allow memory store copying from different a bucket

---
 tests/storage/test_storage.py | 24 ++++++++++++++++++++++++
 xocto/storage/storage.py      |  3 ++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tests/storage/test_storage.py b/tests/storage/test_storage.py
index 5dbc4ed..6f39e10 100644
--- a/tests/storage/test_storage.py
+++ b/tests/storage/test_storage.py
@@ -621,6 +621,30 @@ def test_download_file(self):
         finally:
             os.remove(file.name)
 
+    def test_copy_file(self):
+        # Store a file in a different source bucket
+        source_store = storage.MemoryFileStore("source-bucket")
+        contents = b"test_download_file"
+        source_bucket_name, source_key_path = source_store.store_file(
+            namespace="mem",
+            filename="test.pdf",
+            contents=contents,
+        )
+        s3_object = storage.S3Object(
+            bucket_name=source_bucket_name,
+            key=source_key_path,
+        )
+        destination = "a/b/c.pdf"
+
+        # Copy the file to a different bucket
+        new_s3_object = self.store.copy(s3_object=s3_object, destination=destination)
+
+        # Check the file is copied to the correct bucket and content is as expected
+        assert new_s3_object.bucket_name == self.store.bucket_name
+        assert new_s3_object.key == destination
+        copied_contents = self.store.fetch_file_contents(destination)
+        assert copied_contents == contents
+
 
 class TestLocalFileStore:
     def test_store_and_fetch(self):
diff --git a/xocto/storage/storage.py b/xocto/storage/storage.py
index a45b444..7da9511 100644
--- a/xocto/storage/storage.py
+++ b/xocto/storage/storage.py
@@ -1396,8 +1396,9 @@ def list_s3_keys(self, namespace: str = "") -> Iterable[S3Object]:
         return [self.get_key(path) for path in self.list_files(namespace=namespace)]
 
     def copy(self, *, s3_object: S3Object, destination: str) -> S3Object:
+        source_bucket = self.buffers[s3_object.bucket_name]
         bucket = self.buffers[self.bucket_name]
-        bucket[destination] = bucket[s3_object.key]
+        bucket[destination] = source_bucket[s3_object.key]
         return S3Object(bucket_name=self.bucket_name, key=destination)
 
     def rename(self, *, s3_object: S3Object, destination: str) -> S3Object: