fix(tests): Elasticsearch smoke tests update

entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java

@@ -31,8 +31,10 @@ public class MergedEntityRegistry implements EntityRegistry {
   private final Map<String, AspectSpec> _aspectNameToSpec;
 
   public MergedEntityRegistry(EntityRegistry baseEntityRegistry) {
-    entityNameToSpec = baseEntityRegistry.getEntitySpecs() != null ? baseEntityRegistry.getEntitySpecs() : new HashMap<>();
-    eventNameToSpec = baseEntityRegistry.getEventSpecs() != null ? baseEntityRegistry.getEventSpecs() : new HashMap<>();
+    // baseEntityRegistry.get*Specs() can return immutable Collections.emptyMap() which fails
+    // when this class attempts .put* operations on it.
+    entityNameToSpec = baseEntityRegistry.getEntitySpecs() != null ? new HashMap<>(baseEntityRegistry.getEntitySpecs()) : new HashMap<>();
+    eventNameToSpec = baseEntityRegistry.getEventSpecs() != null ? new HashMap<>(baseEntityRegistry.getEventSpecs()) : new HashMap<>();
     baseEntityRegistry.getAspectTemplateEngine();
     _aspectTemplateEngine = baseEntityRegistry.getAspectTemplateEngine();
     _aspectNameToSpec = baseEntityRegistry.getAspectSpecs();

metadata-io/src/test/java/com/linkedin/metadata/CassandraTestUtils.java

@@ -12,6 +12,7 @@
 import javax.annotation.Nonnull;
 import javax.net.ssl.SSLContext;
 import java.net.InetSocketAddress;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
@@ -34,8 +35,9 @@ public static CassandraContainer setupContainer() {
         .asCompatibleSubstituteFor("cassandra");
 
     CassandraContainer container = new CassandraContainer(imageName);
-    container.withEnv("JVM_OPTS", "-Xms64M -Xmx64M");
-    container.start();
+    container.withEnv("JVM_OPTS", "-Xms64M -Xmx64M")
+            .withStartupTimeout(Duration.ofMinutes(2))
+            .start();
 
     try (Session session = container.getCluster().connect()) {
       session.execute(String.format("CREATE KEYSPACE IF NOT EXISTS %s WITH replication = \n"

metadata-io/src/test/java/com/linkedin/metadata/ElasticSearchTestConfiguration.java

@@ -19,6 +19,7 @@
 
 import javax.annotation.Nonnull;
 
+import java.time.Duration;
 import java.util.Map;
 
 import static com.linkedin.metadata.DockerTestUtils.checkContainerEngine;
@@ -42,7 +43,9 @@ public class ElasticSearchTestConfiguration {
     public RestHighLevelClient getElasticsearchClient() {
         ElasticsearchContainer esContainer = new ElasticsearchContainer(DOCKER_IMAGE_NAME);
         checkContainerEngine(esContainer.getDockerClient());
-        esContainer.start();
+        esContainer.withEnv("ES_JAVA_OPTS", "-Xms64m -Xmx128m -XX:MaxDirectMemorySize=134217728")
+                .withStartupTimeout(Duration.ofMinutes(2))
+                .start();
         return buildRestClient(esContainer);
     }
 

smoke-test/requirements.txt

@@ -5,3 +5,4 @@ tenacity
 -e ../metadata-ingestion[datahub-rest,datahub-kafka,mysql]
 slack-sdk==3.18.1
 aiohttp
+pytest-xdist

smoke-test/smoke-dev.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+set -euxo pipefail
+
+# Runs a basic e2e test. It is not meant to be fully comprehensive,
+# but rather should catch obvious bugs before they make it into prod.
+#
+# Script assumptions:
+#   - The gradle build has already been run.
+#   - Python 3.6+ is installed and in the PATH.
+
+# Log the locally loaded images
+# docker images | grep "datahub-"
+
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+cd "$DIR"
+
+python3 -m venv venv
+source venv/bin/activate
+pip install --upgrade pip wheel setuptools
+pip install -r requirements.txt
+
+echo "DATAHUB_VERSION = ${DATAHUB_VERSION:=acryl-datahub 0.0.0.dev0}"
+DATAHUB_TELEMETRY_ENABLED=false datahub docker quickstart --build-locally --standalone_consumers --dump-logs-on-failure
+
+(cd ..; ./gradlew :smoke-test:yarnInstall)
+
+pytest -n3 -rP --durations=20 -vv --continue-on-collection-errors --junit-xml=junit.smoke.xml $@

smoke-test/smoke.sh

@@ -24,4 +24,4 @@ DATAHUB_TELEMETRY_ENABLED=false datahub docker quickstart --standalone_consumers
 
 (cd ..; ./gradlew :smoke-test:yarnInstall)
 
-pytest -rP --durations=20 -vv --continue-on-collection-errors --junit-xml=junit.smoke.xml
+pytest -n3 -rP --durations=20 -vv --continue-on-collection-errors --junit-xml=junit.smoke.xml

smoke-test/tests/delete/delete_test.py

@@ -48,7 +48,7 @@ def test_setup():
     rollback_url = f"{gms_host}/runs?action=rollback"
     session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": True, "safe": False}))
 
-    sleep(3)
+    sleep(10)
 
     assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
     assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)