testing returning back to create page if create another location is selected - part 2

app/controllers/locations_controller.rb

@@ -10,8 +10,13 @@ class LocationsController < ApplicationController
   def create
     @scheme = Scheme.find(params[:id])
     @location = Location.new(location_params)
-    @location.save
-    redirect_to scheme_check_answers_path(scheme_id: @scheme.id)
+
+    if @location.save
+      redirect_to scheme_check_answers_path(scheme_id: @scheme.id) if location_params[:add_another_location] == "No"
+      redirect_to location_new_scheme_path if location_params[:add_another_location] == "Yes"
+    elsif
+      render :new, status: :unprocessable_entity
+    end
   end
 
 private

spec/requests/locations_controller_spec.rb

@@ -96,7 +96,6 @@ RSpec.describe LocationsController, type: :request do
       end
 
       it "creates a new location for scheme with valid params" do
-
         expect(Location.last.scheme.organisation_id).to eq(user.organisation_id)
         expect(Location.last.name).to eq("Test")
         expect(Location.last.total_units).to eq(5)
@@ -105,14 +104,15 @@ RSpec.describe LocationsController, type: :request do
       end
 
       context "when do you want to add another location is selected as yes" do
+        let(:params) { { location: { name: "Test", total_units: "5", type_of_unit: "Bungalow", wheelchair_adaptation: "No", add_another_location: "Yes" } } }
         it "creates a new location for scheme with valid params and redirects to correct page" do
           expect { post "/schemes/#{scheme.id}/location/create", params: }.to change(Location, :count).by(1)
+          follow_redirect!
           expect(response).to have_http_status(:ok)
           expect(page).to have_content("Add a location to this scheme")
         end
 
         it "creates a new location for scheme with valid params" do
-
           expect(Location.last.scheme.organisation_id).to eq(user.organisation_id)
           expect(Location.last.name).to eq("Test")
           expect(Location.last.total_units).to eq(5)