From 9efb5b59e54b136a6b805da3495874083783d37b Mon Sep 17 00:00:00 2001
From: Rob Winch <rwinch@gopivotal.com>
Date: Mon, 9 Jan 2017 09:10:32 -0600
Subject: [PATCH] Polish custom-cookies

---
 .../java/sample/AttributeTests.java           | 22 ++++--
 .../java/sample/pages/HomePage.java           | 75 +++++++++++--------
 2 files changed, 62 insertions(+), 35 deletions(-)

diff --git a/samples/custom-cookie/src/integration-test/java/sample/AttributeTests.java b/samples/custom-cookie/src/integration-test/java/sample/AttributeTests.java
index 191ab84e..97b8342f 100644
--- a/samples/custom-cookie/src/integration-test/java/sample/AttributeTests.java
+++ b/samples/custom-cookie/src/integration-test/java/sample/AttributeTests.java
@@ -16,17 +16,21 @@
 
 package sample;
 
+import java.util.List;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.htmlunit.HtmlUnitDriver;
 import sample.pages.HomePage;
+import sample.pages.HomePage.Attribute;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * @author Eddú Meléndez
+ * @author Rob Winch
  */
 public class AttributeTests {
 
@@ -51,16 +55,24 @@ public class AttributeTests {
 	@Test
 	public void noAttributes() {
 		HomePage home = HomePage.go(this.driver);
-		assertThat(home.attributes().size()).isEqualTo(0);
+		assertThat(home.attributes()).isEmpty();
 	}
 
 	@Test
 	public void createAttribute() {
 		HomePage home = HomePage.go(this.driver);
-		home.addAttribute("a", "b");
-		assertThat(home.attributes().size()).isEqualTo(1);
-		assertThat(home.row(0).getAttributeName()).isEqualTo("a");
-		assertThat(home.row(0).getAttributeValue()).isEqualTo("b");
+		// @formatter:off
+		home = home.form()
+				.attributeName("a")
+				.attributeValue("b")
+				.submit(HomePage.class);
+		// @formatter:on
+
+		List<Attribute> attributes = home.attributes();
+		assertThat(attributes).hasSize(1);
+		Attribute row = attributes.get(0);
+		assertThat(row.getAttributeName()).isEqualTo("a");
+		assertThat(row.getAttributeValue()).isEqualTo("b");
 	}
 
 }
diff --git a/samples/custom-cookie/src/integration-test/java/sample/pages/HomePage.java b/samples/custom-cookie/src/integration-test/java/sample/pages/HomePage.java
index 7bc855fb..ead76787 100644
--- a/samples/custom-cookie/src/integration-test/java/sample/pages/HomePage.java
+++ b/samples/custom-cookie/src/integration-test/java/sample/pages/HomePage.java
@@ -30,28 +30,23 @@ import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * @author Eddú Meléndez
+ * @author Rob Winch
  */
 public class HomePage {
 
 	private WebDriver driver;
 
-	@FindBy(name = "attributeName")
-	WebElement attributeName;
-
-	@FindBy(name = "attributeValue")
-	WebElement attributeValue;
-
-	@FindBy(css = "input[type=\"submit\"]")
-	WebElement submit;
+	@FindBy(css = "form")
+	WebElement form;
 
 	@FindBy(css = "table tbody tr")
 	List<WebElement> trs;
 
-	private List<Row> rows;
+	List<Attribute> attributes;
 
 	public HomePage(WebDriver driver) {
 		this.driver = driver;
-		this.rows = new ArrayList<Row>();
+		this.attributes = new ArrayList<Attribute>();
 	}
 
 	private static void get(WebDriver driver, String get) {
@@ -68,38 +63,58 @@ public class HomePage {
 		assertThat(this.driver.getTitle()).isEqualTo("Session Attributes");
 	}
 
-	public void addAttribute(String name, String value) {
-		this.attributeName.sendKeys(name);
-		this.attributeValue.sendKeys(value);
-
-		this.submit.click();
-	}
-
-	public List<Row> attributes() {
-		List<Row> rows = new ArrayList<Row>();
+	public List<Attribute> attributes() {
+		List<Attribute> rows = new ArrayList<Attribute>();
 		for (WebElement tr : this.trs) {
-			rows.add(new Row(tr));
+			rows.add(new Attribute(tr));
 		}
-		this.rows.addAll(rows);
-		return this.rows;
+		this.attributes.addAll(rows);
+		return this.attributes;
 	}
 
-	public Row row(int index) {
-		return this.rows.get(index);
+	public Form form() {
+		return new Form(this.form);
 	}
 
-	public static class Row {
+	public class Form {
+		@FindBy(name = "attributeName")
+		WebElement attributeName;
+
+		@FindBy(name = "attributeValue")
+		WebElement attributeValue;
+
+		@FindBy(css = "input[type=\"submit\"]")
+		WebElement submit;
+
+		public Form(SearchContext context) {
+			PageFactory.initElements(new DefaultElementLocatorFactory(context), this);
+		}
+
+		public Form attributeName(String text) {
+			this.attributeName.sendKeys(text);
+			return this;
+		}
+
+		public Form attributeValue(String text) {
+			this.attributeValue.sendKeys(text);
+			return this;
+		}
+
+		public <T> T submit(Class<T> page) {
+			this.submit.click();
+			return PageFactory.initElements(HomePage.this.driver, page);
+		}
+	}
+
+	public static class Attribute {
 		@FindBy(xpath = "//td[1]")
 		WebElement attributeName;
 
 		@FindBy(xpath = "//td[2]")
 		WebElement attributeValue;
 
-		public Row(SearchContext context) {
-			super();
-			DefaultElementLocatorFactory factory = new DefaultElementLocatorFactory(
-					context);
-			PageFactory.initElements(factory, this);
+		public Attribute(SearchContext context) {
+			PageFactory.initElements(new DefaultElementLocatorFactory(context), this);
 		}
 
 		/**