Add incomplete tests for startCluster

- add test for generate_ecs_config
- add test for checking spot fleet config file
- modify example fleet file KeyName to match config.py as stated in docs

Author: gnodar01

files/exampleFleet_us-east-1.json

@@ -5,7 +5,7 @@
   "LaunchSpecifications": [
     {
       "ImageId": "ami-fad25980",
-      "KeyName": "your_key_file_name",
+      "KeyName": "your-key-file",
       "IamInstanceProfile": {
         "Arn": "arn:aws:iam::XXXXXXXXXXXX:instance-profile/ecsInstanceRole"
       },

files/exampleFleet_us-west-2.json

@@ -5,7 +5,7 @@
   "LaunchSpecifications": [
     {
       "ImageId": "ami-c9c87cb1",
-      "KeyName": "your_key_file_name",
+      "KeyName": "your-key-file",
       "IamInstanceProfile": {
         "Arn": "arn:aws:iam::XXXXXXXXXXXX:instance-profile/ecsInstanceRole"
       },

tests/conftest.py

@@ -4,10 +4,10 @@
 
 import pytest
 import boto3
-from moto import mock_sqs, mock_ecs
+from moto import mock_sqs, mock_ecs, mock_s3
 
 import run
-from config import AWS_REGION, AWS_PROFILE
+from config import AWS_REGION, AWS_PROFILE, AWS_BUCKET, APP_NAME
 
 
 # WARNING: Do not import a module here or in any of the tests
@@ -32,6 +32,8 @@
 
 JOB_FILE = Path(__file__).parent.parent / "files/exampleJob.json"
 FLEET_FILE = Path(__file__).parent.parent / "files/exampleFleet_us-east-1.json"
+# does not exist unless startCluster has been run
+MONITOR_FILE = Path(__file__).parent.parent / f"files/{APP_NAME}SpotFleetRequestId.json"
 
 
 @pytest.fixture(autouse=True, scope="session")
@@ -95,13 +97,34 @@ def aws_config(monkeypatch, tmp_path):
 @pytest.fixture(scope="function")
 def sqs():
     with mock_sqs():
-        yield boto3.client("sqs", region_name=os.environ["AWS_DEFAULT_REGION"])
+        yield boto3.client("sqs", region_name=AWS_REGION)
 
 
 @pytest.fixture(scope="function")
 def ecs():
     with mock_ecs():
-        yield boto3.client("ecs", os.environ["AWS_DEFAULT_REGION"])
+        yield boto3.client("ecs", region_name=AWS_REGION)
+
+
+@pytest.fixture(scope="function")
+def s3():
+    with mock_s3():
+        yield boto3.client("s3", region_name=AWS_REGION)
+
+
+@pytest.fixture(scope="function")
+def protect_monitor_file():
+    # read in the current contents, if any
+    curr = None
+    if MONITOR_FILE.exists():
+        curr = MONITOR_FILE.read_text()
+
+    # do something that may or may not write to the file
+    yield MONITOR_FILE
+    
+    # clean up by putting the original contents back, if any
+    if curr:
+        MONITOR_FILE.write_text(curr)
 
 
 # Below functions are fixtures that run steps 1 - 4
@@ -116,6 +139,8 @@ def ecs():
 # Therefor we return a non-mocked callback, the tests are decorated with
 # mocks, and then the tests invoke the callback returned by the fixture.
 
+
+# mock sqs and ecs before running cb
 @pytest.fixture(scope="function")
 def run_setup(aws_config):
     def f():
@@ -124,12 +149,38 @@ def f():
     return f
 
 
+# mock sqs and ecs before running cb
 @pytest.fixture(scope="function")
 def run_submitJob(run_setup, monkeypatch):
-    monkeypatch.setattr(sys, "argv", ["run.py", "submitJob", str(JOB_FILE)])
-
     def f():
+        # don't put this outside of the callback, else it may be overwritten
+        monkeypatch.setattr(sys, "argv", ["run.py", "submitJob", str(JOB_FILE)])
+
         run_setup()
         run.submitJob()
+        print('x')
+
+    return f
+
+
+# mock sqs, ecs, s3 and ec2 before running cb
+@pytest.fixture(scope="function")
+def run_startCluster(run_submitJob, monkeypatch, protect_monitor_file):
+    def f():
+        s3 = boto3.client('s3')
+
+        if (AWS_REGION == "us-east-1"):
+            # 'us-east-1' is the default region for S3 buckets
+            # and is not a vallid arg for "LocationConstraint"
+            s3.create_bucket(Bucket=AWS_BUCKET)
+        else:
+            s3.create_bucket(Bucket=AWS_BUCKET, CreateBucketConfiguration={"LocationConstraint": AWS_REGION})
+
+        run_submitJob()
+        
+        # don't put this outside of the callback, else it may be overwritten
+        monkeypatch.setattr(sys, "argv", ["run.py", "startCluster", str(FLEET_FILE)])
+        
+        run.startCluster()
 
     return f

tests/test_startCluster.py

@@ -0,0 +1,143 @@
+import pytest
+from moto import mock_ecs, mock_sqs, mock_s3, mock_ec2
+import boto3
+
+import run
+import config
+
+
+# startCluster will:
+#   send a spot fleet request to moto-mocked-aws
+#   create a APP_NAMESpotFleetRequestId.json file
+#   once spot fleet request is ready (instantly), creat log groups for your log streams to go in
+#   DS will ask aws moto-mocked-aws to place Docker containers onto the spot fleet instances
+#   job will begin instantly
+
+class TestGenerateECSConfig:
+    @mock_ecs
+    @mock_sqs
+    def test_generate_ecs_config(self, s3, run_submitJob, tmp_path):
+        if (config.AWS_REGION == "us-east-1"):
+            # 'us-east-1' is the default region for S3 buckets
+            # and is not a vallid arg for "LocationConstraint"
+            s3.create_bucket(Bucket=config.AWS_BUCKET)
+        else:
+            s3.create_bucket(Bucket=config.AWS_BUCKET, CreateBucketConfiguration={"LocationConstraint": config.AWS_REGION})
+
+        run_submitJob()
+
+        res_endpoint = run.generateECSconfig(config.ECS_CLUSTER, config.APP_NAME, config.AWS_BUCKET, s3)
+
+        expected_key = f"ecsconfigs/{config.APP_NAME}_ecs.config"
+        expected_file_path = tmp_path / "configtemp.config"
+        expected_file_endpoint = f"s3://{config.AWS_BUCKET}/{expected_key}"
+
+        assert res_endpoint == expected_file_endpoint
+
+        # create the file object to write to
+        expected_file_path.touch()
+
+        with expected_file_path.open('wb') as f:
+            s3.download_fileobj(config.AWS_BUCKET, expected_key, f)
+        
+        res_file = expected_file_path.read_text()
+
+        assert res_file == f"ECS_CLUSTER={config.ECS_CLUSTER}\nECS_AVAILABLE_LOGGING_DRIVERS=[\"json-file\",\"awslogs\"]"
+
+
+
+class EarlyTermination(Exception):
+    ...
+
+def hijack_ec2(real_client):
+    """
+    Patches boto3.client so that an invocation of 'ec2' will cause an
+    EarlyTermination exception, allowing inspection and testing of the
+    stack frame up until that point.
+    """
+    def f(*args, **kwargs):
+        if (args[0] == 'ec2'):
+            raise EarlyTermination("early termination")
+        
+        return real_client(*args, **kwargs)
+    
+    return f
+
+
+class TestSpotFleetConfig:
+    @mock_ecs
+    @mock_sqs
+    @mock_s3
+    def test_spot_fleet_config(self, run_startCluster, monkeypatch):
+        monkeypatch.setattr(boto3, "client", hijack_ec2(boto3.client))
+        with pytest.raises(EarlyTermination) as e_info:
+            run_startCluster()
+
+        spot_fleet_config_res = None
+        for tb in e_info.traceback:
+            if (tb.name == "startCluster"):
+                spot_fleet_config_res = tb.frame.f_locals["spotfleetConfig"]
+        
+        assert spot_fleet_config_res is not None
+
+        # For config file requirements, see:
+        # https://distributedscience.github.io/Distributed-Something/step_3_start_cluster.html#configuring-your-spot-fleet-request
+        
+        assert "IamFleetRole" in spot_fleet_config_res
+        assert spot_fleet_config_res["IamFleetRole"].startswith("arn:aws:iam::")
+
+        assert "ValidFrom" in spot_fleet_config_res
+        assert "ValidUntil" in spot_fleet_config_res
+        assert "TargetCapacity" in spot_fleet_config_res
+        assert spot_fleet_config_res["TargetCapacity"] == config.CLUSTER_MACHINES
+        assert "SpotPrice" in spot_fleet_config_res
+        assert float(spot_fleet_config_res["SpotPrice"]) == pytest.approx(config.MACHINE_PRICE, rel=1e-2)
+
+        launch_specs = spot_fleet_config_res["LaunchSpecifications"]
+        for i in range(len(launch_specs)):
+            assert "IamInstanceProfile" in launch_specs[i]
+            assert "Arn" in launch_specs[i]["IamInstanceProfile"]
+            assert launch_specs[i]["IamInstanceProfile"]["Arn"].startswith("arn:aws:iam::")
+
+            assert "KeyName" in launch_specs[i]
+            assert launch_specs[i]["KeyName"] == config.SSH_KEY_NAME[:-4]
+
+            assert "ImageId" in launch_specs[i]
+            assert launch_specs[i]["ImageId"].startswith("ami-")
+
+            assert "NetworkInterfaces" in launch_specs[i]
+            net_intfcs = launch_specs[i]["NetworkInterfaces"]
+            for j in range(len(net_intfcs)):
+                assert "SubnetId" in net_intfcs[j]
+                assert net_intfcs[j]["SubnetId"].startswith("subnet-")
+
+                assert "Groups" in net_intfcs[j]
+                grps = net_intfcs[j]["Groups"]
+                for k in range(len(grps)):
+                    assert grps[k].startswith("sg-")
+
+            assert "BlockDeviceMappings" in launch_specs[i]
+            bdms = launch_specs[i]["BlockDeviceMappings"]
+            
+            assert "Ebs" in bdms[0]
+            assert "SnapshotId" in bdms[0]["Ebs"]
+            assert bdms[0]["Ebs"]["SnapshotId"].startswith("snap-")
+
+            assert "Ebs" in bdms[1]
+            assert "VolumeSize" in bdms[1]["Ebs"]
+            assert bdms[1]["Ebs"]["VolumeSize"] == config.EBS_VOL_SIZE
+
+            assert "InstanceType" in launch_specs[i]
+            assert launch_specs[i]["InstanceType"] == config.MACHINE_TYPE[i]
+
+            assert "UserData" in launch_specs[i]
+
+
+class TestStartCluster:
+    @mock_ecs
+    @mock_sqs
+    @mock_s3
+    @mock_ec2
+    @pytest.mark.skip(reason="not implemented yet")
+    def test_start_cluster(self, run_startCluster):
+        run_startCluster()