From 3bec2a480f8d17e8e22cbd23fac2b4611a786e36 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@itamarst.org>
Date: Wed, 8 Sep 2021 12:20:27 -0400
Subject: [PATCH] Start on allocate_bucket tests.

---
 src/allmydata/test/test_istorageserver.py | 152 +++++++++++++++++++++-
 1 file changed, 149 insertions(+), 3 deletions(-)

diff --git a/src/allmydata/test/test_istorageserver.py b/src/allmydata/test/test_istorageserver.py
index b2229f467..303b4b997 100644
--- a/src/allmydata/test/test_istorageserver.py
+++ b/src/allmydata/test/test_istorageserver.py
@@ -6,13 +6,68 @@ reused across tests, so each test should be careful to generate unique storage
 indexes.
 """
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from future.utils import PY2
+
+if PY2:
+    from future.builtins import (
+        filter,
+        map,
+        zip,
+        ascii,
+        chr,
+        hex,
+        input,
+        next,
+        oct,
+        open,
+        pow,
+        round,
+        super,
+        bytes,
+        dict,
+        list,
+        object,
+        range,
+        str,
+        max,
+        min,
+    )  # noqa: F401
+
+from random import randrange
+from unittest import expectedFailure
+
 from twisted.internet.defer import inlineCallbacks, returnValue
 from twisted.trial.unittest import TestCase
 
-from allmydata.interfaces import IStorageServer
+from foolscap.api import Referenceable
+
+from allmydata.interfaces import IStorageServer, RIBucketWriter
 from .test_system import SystemTestMixin
 
 
+def _randbytes(length):
+    # type: (int) -> bytes
+    """Return random bytes string of given length."""
+    return bytes([randrange(0, 256) for _ in range(length)])
+
+
+def new_storage_index():
+    # type: () -> bytes
+    """Return a new random storage index."""
+    return _randbytes(16)
+
+
+def new_secret():
+    # type: () -> bytes
+    """Return a new random secret (for lease renewal or cancellation)."""
+    return _randbytes(32)
+
+
 class IStorageServerSharedAPIsTestsMixin(object):
     """
     Tests for ``IStorageServer``'s shared APIs.
@@ -35,13 +90,104 @@ class IStorageServerImmutableAPIsTestsMixin(object):
     """
 
     # TODO === allocate_buckets + RIBucketWriter ===
-    # TODO allocate_buckets on a new storage index
-    # TODO allocate_buckets on existing bucket with same sharenums
+    # DONE allocate_buckets on a new storage index
+    # PROG allocate_buckets on existing bucket with same sharenums
     # TODO allocate_buckets with smaller sharenums
     # TODO allocate_buckets with larger sharenums
     # TODO writes to bucket can happen in any order (write then read)
     # TODO overlapping writes ignore already-written data (write then read)
 
+    @inlineCallbacks
+    def test_allocate_buckets_new(self):
+        """
+        allocate_buckets() with a new storage index returns the matching
+        shares.
+        """
+        (already_got, allocated) = yield self.storage_server.allocate_buckets(
+            new_storage_index(),
+            new_secret(),
+            new_secret(),
+            set(range(5)),
+            1024,
+            Referenceable(),
+        )
+        self.assertEqual(already_got, set())
+        self.assertEqual(allocated.keys(), set(range(5)))
+        # We validate the bucket objects' interface in a later test.
+
+    @inlineCallbacks
+    def test_allocate_buckets_repeat(self):
+        """
+        allocate_buckets() with the same storage index returns the same result,
+        because the shares have not been written to.
+
+        This fails due to https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3793
+        """
+        si, renew_secret, cancel_secret = (
+            new_storage_index(),
+            new_secret(),
+            new_secret(),
+        )
+        (already_got, allocated) = yield self.storage_server.allocate_buckets(
+            si,
+            renew_secret,
+            cancel_secret,
+            set(range(5)),
+            1024,
+            Referenceable(),
+        )
+        (already_got2, allocated2) = yield self.storage_server.allocate_buckets(
+            si,
+            renew_secret,
+            cancel_secret,
+            set(range(5)),
+            1024,
+            Referenceable(),
+        )
+        self.assertEqual(already_got, already_got2)
+        self.assertEqual(allocated.keys(), allocated2.keys())
+
+    test_allocate_buckets_repeat.todo = (
+        "https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3793"
+    )
+
+    @expectedFailure
+    @inlineCallbacks
+    def test_allocate_buckets_more_sharenums(self):
+        """
+        allocate_buckets() with the same storage index but more sharenums
+        acknowledges the extra shares don't exist.
+
+        Fails due to https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3793
+        """
+        si, renew_secret, cancel_secret = (
+            new_storage_index(),
+            new_secret(),
+            new_secret(),
+        )
+        yield self.storage_server.allocate_buckets(
+            si,
+            renew_secret,
+            cancel_secret,
+            set(range(5)),
+            1024,
+            Referenceable(),
+        )
+        (already_got2, allocated2) = yield self.storage_server.allocate_buckets(
+            si,
+            renew_secret,
+            cancel_secret,
+            set(range(7)),
+            1024,
+            Referenceable(),
+        )
+        self.assertEqual(already_got2, set())  # none were fully written
+        self.assertEqual(allocated2.keys(), set(range(7)))
+
+    test_allocate_buckets_more_sharenums.todo = (
+        "https://tahoe-lafs.org/trac/tahoe-lafs/ticket/3793"
+    )
+
 
 class _FoolscapMixin(SystemTestMixin):
     """Run tests on Foolscap version of ``IStorageServer."""