fix tests

bzknightz · Sep 1, 2021 · 7898842 · 7898842
1 parent 7820031
commit 7898842
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 13 deletions.
diff --git a/ci/tests/cosmosdb/run-cosmosdb-server.sh b/ci/tests/cosmosdb/run-cosmosdb-server.sh
@@ -4,8 +4,8 @@ ipaddr="$(ifconfig | grep "inet " | grep -Fv 127.0.0.1 | awk '{print $2}' | head
 echo "System IP address is $ipaddr"
 docker pull mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator
 docker run --rm -p 8081:8081 -p 10251:10251 -p 10252:10252 -p 10253:10253 -p 10254:10254 \
-  -m 3g --cpus=2.0 --name=cosmosdb \
-  -e AZURE_COSMOS_EMULATOR_PARTITION_COUNT=2 \
+  -m 4g --cpus=2.0 --name=cosmosdb \
+  -e AZURE_COSMOS_EMULATOR_PARTITION_COUNT=4 \
   -e AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE=true \
   -e AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE="$ipaddr" \
   -d mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator

diff --git a/...-server-support-cosmosdb-core/src/main/java/org/apereo/cas/cosmosdb/CosmosDbDocument.java b/...-server-support-cosmosdb-core/src/main/java/org/apereo/cas/cosmosdb/CosmosDbDocument.java
@@ -18,5 +18,7 @@ public class CosmosDbDocument implements Serializable {
 
     private String id;
 
+    private String serviceId;
+
     private String body;
 }
diff --git a/support/cas-server-support-cosmosdb-service-registry/build.gradle b/support/cas-server-support-cosmosdb-service-registry/build.gradle
@@ -1,4 +1,7 @@
 description = "Apereo CAS Azure Cosmos DB Service Registry"
+ext {
+    maxParallelForksForTests = 1
+}
 dependencies {
     api project(":api:cas-server-core-api-services")
 

diff --git a/...mosdb-service-registry/src/main/java/org/apereo/cas/services/CosmosDbServiceRegistry.java b/...mosdb-service-registry/src/main/java/org/apereo/cas/services/CosmosDbServiceRegistry.java
@@ -3,6 +3,7 @@
 import org.apereo.cas.cosmosdb.CosmosDbDocument;
 import org.apereo.cas.cosmosdb.CosmosDbObjectFactory;
 import org.apereo.cas.services.util.RegisteredServiceJsonSerializer;
+import org.apereo.cas.util.LoggingUtils;
 import org.apereo.cas.util.serialization.StringSerializer;
 
 import com.azure.cosmos.CosmosContainer;
@@ -33,8 +34,6 @@ public class CosmosDbServiceRegistry extends AbstractServiceRegistry {
      */
     public static final String PARTITION_KEY = "id";
 
-    private final CosmosDbObjectFactory cosmosDbFactory;
-
     private final StringSerializer<RegisteredService> serializer;
 
     private final CosmosContainer container;
@@ -44,9 +43,8 @@ public CosmosDbServiceRegistry(final CosmosDbObjectFactory factory,
                                    final ConfigurableApplicationContext applicationContext,
                                    final Collection<ServiceRegistryListener> serviceRegistryListeners) {
         super(applicationContext, serviceRegistryListeners);
-        this.cosmosDbFactory = factory;
         this.serializer = new RegisteredServiceJsonSerializer(new MinimalPrettyPrinter());
-        cosmosDbFactory.createContainer(containerName, PARTITION_KEY);
+        factory.createContainer(containerName, PARTITION_KEY);
         this.container = factory.getContainer(containerName);
     }
 
@@ -65,13 +63,14 @@ public RegisteredService save(final RegisteredService registeredService) {
     @Override
     public boolean delete(final RegisteredService registeredService) {
         val doc = createCosmosDbDocument(registeredService);
-        LOGGER.debug("Loading registered services [{}] from container [{}]", doc.getId(), container.getId());
+        LOGGER.debug("Deleting registered service [{}] from container [{}]", doc.getId(), container.getId());
         val response = container.deleteItem(doc, new CosmosItemRequestOptions());
         return !HttpStatus.valueOf(response.getStatusCode()).isError();
     }
 
     @Override
     public void deleteAll() {
+        LOGGER.debug("Deleting registered services from container [{}]", container.getId());
         val queryOptions = new CosmosQueryRequestOptions();
         val items = container.queryItems("SELECT * FROM " + container.getId(), queryOptions, CosmosDbDocument.class);
         items.iterableByPage()
@@ -83,7 +82,7 @@ public void deleteAll() {
     public Collection<RegisteredService> load() {
         val services = new ArrayList<RegisteredService>();
         val queryOptions = new CosmosQueryRequestOptions();
-        LOGGER.debug("Loading registered services from container [{}]", container.getId());
+        LOGGER.trace("Loading registered services from container [{}]", container.getId());
         val items = container.queryItems("SELECT * FROM " + container.getId(), queryOptions, CosmosDbDocument.class);
         items.iterableByPage()
             .forEach(response -> services.addAll(response.getResults()
@@ -94,6 +93,24 @@ public Collection<RegisteredService> load() {
         return services;
     }
 
+    @Override
+    public RegisteredService findServiceBy(final String id) {
+        val services = new ArrayList<RegisteredService>();
+        val queryOptions = new CosmosQueryRequestOptions();
+        LOGGER.trace("Loading registered service by [{}] from container [{}]", id, container.getId());
+        val query = String.format("SELECT * FROM %s r WHERE r.serviceId LIKE '%s'", container.getId(), '%' + id + '%');
+        val items = container.queryItems(query, queryOptions, CosmosDbDocument.class);
+        items.iterableByPage()
+            .forEach(response -> services.addAll(response.getResults()
+                .stream()
+                .map(this::getRegisteredServiceFromDocumentBody)
+                .sorted()
+                .filter(r -> r.matches(id))
+                .peek(this::invokeServiceRegistryListenerPostLoad)
+                .collect(Collectors.toList())));
+        return services.isEmpty() ? null : services.get(0);
+    }
+
     @Override
     public RegisteredService findServiceById(final long id) {
         try {
@@ -102,7 +119,12 @@ public RegisteredService findServiceById(final long id) {
             val doc = container.readItem(key, new PartitionKey(key), CosmosDbDocument.class).getItem();
             return getRegisteredServiceFromDocumentBody(doc);
         } catch (final CosmosException e) {
-            LOGGER.debug(e.getMessage(), e);
+            if (e.getStatusCode() == HttpStatus.NOT_FOUND.value()) {
+                LOGGER.debug("Unable to locate registered service with id [{}]", id);
+                LOGGER.trace(e.getMessage(), e);
+            } else {
+                LoggingUtils.warn(LOGGER, e);
+            }
         }
         return null;
     }
@@ -113,13 +135,13 @@ private RegisteredService getRegisteredServiceFromDocumentBody(final CosmosDbDoc
 
     private void insert(final RegisteredService registeredService) {
         val doc = createCosmosDbDocument(registeredService);
-        LOGGER.debug("Creating registered service [{}] in container [{}]", doc.getId(), container.getId());
+        LOGGER.trace("Creating registered service with id [{}] in container [{}]", doc.getId(), container.getId());
         container.createItem(doc);
     }
 
     private void update(final RegisteredService registeredService) {
         val doc = createCosmosDbDocument(registeredService);
-        LOGGER.debug("Upserting registered service [{}] in container [{}]", doc.getId(), container.getId());
+        LOGGER.trace("Upserting registered service with id [{}] in container [{}]", doc.getId(), container.getId());
         container.upsertItem(doc);
     }
 
@@ -128,6 +150,7 @@ private CosmosDbDocument createCosmosDbDocument(final RegisteredService register
         val document = new CosmosDbDocument();
         document.setBody(body);
         document.setId(String.valueOf(registeredService.getId()));
+        document.setServiceId(registeredService.getServiceId());
         return document;
     }
 }
diff --git a/...-service-registry/src/test/java/org/apereo/cas/services/CosmosDbServiceRegistryTests.java b/...-service-registry/src/test/java/org/apereo/cas/services/CosmosDbServiceRegistryTests.java
@@ -43,7 +43,7 @@
     "cas.service-registry.cosmosDb.uri=https://localhost:8081",
     "cas.service-registry.cosmosDb.key=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
     "cas.service-registry.cosmosDb.database=RegisteredServicesDb",
-    "cas.service-registry.cosmosDb.max-retry-attempts-on-throttled-requests=1",
+    "cas.service-registry.cosmosDb.max-retry-attempts-on-throttled-requests=5",
     "cas.service-registry.cosmosDb.indexing-mode=CONSISTENT",
     "cas.service-registry.cosmosDb.drop-container=true"
 })
@@ -57,7 +57,7 @@ public class CosmosDbServiceRegistryTests extends AbstractServiceRegistryTests {
 
     @BeforeEach
     public void deleteAll() {
-        newServiceRegistry.load().forEach(service -> newServiceRegistry.delete(service));
+        newServiceRegistry.deleteAll();
         assertTrue(newServiceRegistry.load().isEmpty());
     }
 }