Upgradeable PriceEstimator with test and deploy script

contracts/PriceEstimator/IPriceEstimator.sol

@@ -0,0 +1,14 @@
+pragma solidity ^0.4.18;
+
+import "./../Upgradeability/OwnableUpgradeableImplementation/IOwnableUpgradeableImplementation.sol";
+
+contract IPriceEstimator is IOwnableUpgradeableImplementation {
+
+    function setHouseholdMetersContract(address _householdMetersAddress) public returns(bool success);
+
+    function setWaterVouchersContract(address _waterVoucherAddress) public returns(bool success);
+
+    function getCurrentMonthLiters(address _meter) public view returns(uint256 liters);
+
+    function estimate(address _meter, uint _liters) public view returns(uint256 amount, uint256 price); 
+}

contracts/PriceEstimator/PriceEstimator.sol

@@ -2,17 +2,23 @@ pragma solidity ^0.4.18;
 
 import "./../HouseholdMeters.sol";
 import "./../WaterVouchers.sol";
+import "./../Upgradeability/OwnableUpgradeableImplementation/OwnableUpgradeableImplementation.sol";
+import "./IPriceEstimator.sol";
 import "zeppelin-solidity/contracts/math/SafeMath.sol";
 
-contract PriceEstimator {
+contract PriceEstimator is IPriceEstimator, OwnableUpgradeableImplementation {
     using SafeMath for uint256;
-// TODO Make it upgradeble
     HouseholdMeters householdMetersContract;
     WaterVouchers waterVouchersContract;
 
-    function PriceEstimator (address _householdMetersAddress, address _waterVoucherAddress) public {
+    function setHouseholdMetersContract(address _householdMetersAddress) public onlyOwner returns(bool success) {
         householdMetersContract = HouseholdMeters(_householdMetersAddress);
+        return true;
+    }
+
+    function setWaterVouchersContract(address _waterVoucherAddress) public onlyOwner returns(bool success) {
         waterVouchersContract = WaterVouchers(_waterVoucherAddress);
+        return true;
     }
 
     function getCurrentMonthLiters(address _meter) public view returns(uint256 liters) {

contracts/PriceEstimator/PriceEstimatorProxy.sol

@@ -0,0 +1,7 @@
+pragma solidity ^0.4.18;
+
+import "./../Upgradeability/UpgradeableProxy.sol";
+
+contract PriceEstimatorProxy is UpgradeableProxy {
+	function PriceEstimatorProxy(address initialImplementation) public UpgradeableProxy(initialImplementation) {}
+}

migrations/3_waterSystem.js

@@ -2,6 +2,8 @@ let WaterGoverning = artifacts.require("./WaterGoverning.sol");
 let WaterVouchers = artifacts.require("./WaterVouchers.sol");
 var HouseholdMeters = artifacts.require("./HouseholdMeters.sol");
 let PriceEstimator = artifacts.require("./PriceEstimator/PriceEstimator.sol");
+let IPriceEstimator = artifacts.require("./PriceEstimator/IPriceEstimator.sol");
+let PriceEstimatorProxy = artifacts.require("./PriceEstimator/PriceEstimatorProxy.sol");
 
 module.exports = async function (deployer) {
     await deployer.deploy(HouseholdMeters);
@@ -13,11 +15,17 @@ module.exports = async function (deployer) {
     await deployer.deploy(WaterGoverning, WaterVouchersInstance.address);
     let WaterGoverningInstance = await WaterGoverning.deployed();
 
-    await deployer.deploy(PriceEstimator, HouseholdMetersInstance.address, WaterVouchersInstance.address);
-    let PriceEstimatorInstance = await PriceEstimator.deployed();
+    await deployer.deploy(PriceEstimator);
+    let PriceEstimatorImpl = await PriceEstimator.deployed();
+    await deployer.deploy(PriceEstimatorProxy, PriceEstimatorImpl.address);
+    let PriceEstimatorContract = await PriceEstimatorProxy.deployed();
+    PriceEstimatorContract = await IPriceEstimator.at(PriceEstimatorContract.address);
 
+    await PriceEstimatorContract.init();
+    await PriceEstimatorContract.setHouseholdMetersContract(HouseholdMetersInstance.address);
+    await PriceEstimatorContract.setWaterVouchersContract(WaterVouchersInstance.address);
     await WaterVouchersInstance.setWaterGoverningContractAddress(WaterGoverningInstance.address);
-    await WaterVouchersInstance.setPriceEstimatorContractAddress(PriceEstimatorInstance.address);
+    await WaterVouchersInstance.setPriceEstimatorContractAddress(PriceEstimatorContract.address);
 
     // Ganache first Account - for demo purposes - DO NOT USE IT IN PRODUCTION
     await WaterVouchersInstance.addIntermediary("0x627306090abaB3A6e1400e9345bC60c78a8BEf57");

test/PriceEstimator.js

@@ -1,5 +1,8 @@
 const HouseholdMeters = artifacts.require("./HouseholdMeters.sol");
 const PriceEstimator = artifacts.require("./PriceEstimator/PriceEstimator.sol");
+const IPriceEstimator = artifacts.require("./PriceEstimator/IPriceEstimator.sol");
+const PriceEstimatorProxy = artifacts.require("./PriceEstimator/PriceEstimatorProxy.sol");
+const IOwnableUpgradeableImplementation = artifacts.require("./Upgradeability/OwnableUpgradeableImplementation/IOwnableUpgradeableImplementation.sol");
 const WaterVoucher = artifacts.require("./WaterVouchers.sol");
 const WaterGoverning = artifacts.require("./WaterGoverning.sol");
 const util = require('./util');
@@ -35,9 +38,120 @@ contract('PriceEstimator', function (accounts) {
     let HouseholdMetersContract;
     let PriceEstimatorContract;
     let WaterVoucherContract;
+    let PriceEstimatorContractImpl;
+    let PriceEstimatorContractProxy;
 
     var wait = ms => new Promise((r, j) => setTimeout(r, ms))
 
+    describe("Creating PriceEstimator contract", () => {
+        beforeEach(async function () {
+            HouseholdMetersContract = await HouseholdMeters.new({
+                from: _owner
+            });
+
+            WaterVoucherContract = await WaterVoucher.new({
+                from: _owner
+            });
+
+            WaterGoverningContract = await WaterGoverning.new(WaterVoucherContract.address, {
+                from: _owner
+            });
+
+            let PriceEstimatorContractImpl = await PriceEstimator.new({
+                from: _owner
+            });
+
+            let PriceEstimatorContractProxy = await PriceEstimator.new(PriceEstimatorContractImpl.address, {
+                from: _owner
+            });
+
+            PriceEstimatorContract = await IPriceEstimator.at(PriceEstimatorContractProxy.address);
+
+            await PriceEstimatorContract.init({
+                from: _owner
+            });
+            await PriceEstimatorContract.setHouseholdMetersContract(HouseholdMetersContract.address);
+            await PriceEstimatorContract.setWaterVouchersContract(WaterVoucherContract.address);
+
+            await WaterVoucherContract.setPriceEstimatorContractAddress(PriceEstimatorContract.address);
+            await WaterVoucherContract.setWaterGoverningContractAddress(WaterGoverningContract.address);
+
+            await HouseholdMetersContract.addAdmin(_admin1);
+            await HouseholdMetersContract.addHouseholdMeter(_householdMeter, _householdMembers, {
+                from: _admin1
+            });
+
+            await WaterVoucherContract.addIntermediary(_admin1);
+        });
+
+        it("should get the owner of the first contract", async function () {
+            const owner = await PriceEstimatorContract.getOwner();
+            assert.strictEqual(owner, _owner, "The owner is not set correctly");
+        });
+    });
+
+    describe("Upgrade PriceEstimator contract", () => {
+        beforeEach(async function () {
+            HouseholdMetersContract = await HouseholdMeters.new({
+                from: _owner
+            });
+
+            WaterVoucherContract = await WaterVoucher.new({
+                from: _owner
+            });
+
+            WaterGoverningContract = await WaterGoverning.new(WaterVoucherContract.address, {
+                from: _owner
+            });
+
+            PriceEstimatorContractImpl = await PriceEstimator.new({
+                from: _owner
+            });
+
+            PriceEstimatorContractProxy = await PriceEstimator.new(PriceEstimatorContractImpl.address, {
+                from: _owner
+            });
+
+            PriceEstimatorContract = await IPriceEstimator.at(PriceEstimatorContractProxy.address);
+
+            await PriceEstimatorContract.init({
+                from: _owner
+            });
+            await PriceEstimatorContract.setHouseholdMetersContract(HouseholdMetersContract.address);
+            await PriceEstimatorContract.setWaterVouchersContract(WaterVoucherContract.address);
+
+            await WaterVoucherContract.setPriceEstimatorContractAddress(PriceEstimatorContract.address);
+            await WaterVoucherContract.setWaterGoverningContractAddress(WaterGoverningContract.address);
+
+            await HouseholdMetersContract.addAdmin(_admin1);
+            await HouseholdMetersContract.addHouseholdMeter(_householdMeter, _householdMembers, {
+                from: _admin1
+            });
+
+            await WaterVoucherContract.addIntermediary(_admin1);
+        });
+
+        it("should upgrade contract from owner", async function () {
+            let PriceEstimatorContractImpl2 = await PriceEstimator.new({
+                from: _owner
+            });
+            const upgradeableContract = await IOwnableUpgradeableImplementation.at(PriceEstimatorContractProxy.address);
+            await upgradeableContract.upgradeImplementation(PriceEstimatorContractImpl2.address);
+            const newImplAddress = await upgradeableContract.getImplementation();
+            assert.strictEqual(PriceEstimatorContractImpl2.address, newImplAddress, "The address is not set correctly");
+        });
+
+        it("should throw on upgrade contract from not owner", async function () {
+            let PriceEstimatorContractImpl2 = await PriceEstimator.new({
+                from: _owner
+            });
+            const upgradeableContract = await IOwnableUpgradeableImplementation.at(PriceEstimatorContractProxy.address);
+            await expectThrow(upgradeableContract.upgradeImplementation(PriceEstimatorContractImpl2.address, {
+                from: _notOwner
+            }));
+        });
+    });
+
     describe("Estimate correct prices", () => {
         beforeEach(async function () {
             HouseholdMetersContract = await HouseholdMeters.new({
@@ -52,12 +166,24 @@ contract('PriceEstimator', function (accounts) {
                 from: _owner
             });
 
-            PriceEstimatorContract = await PriceEstimator.new(HouseholdMetersContract.address, WaterVoucherContract.address, {
+            let PriceEstimatorContractImpl = await PriceEstimator.new({
+                from: _owner
+            });
+
+            let PriceEstimatorContractProxy = await PriceEstimator.new(PriceEstimatorContractImpl.address, {
                 from: _owner
             });
 
-            WaterVoucherContract.setPriceEstimatorContractAddress(PriceEstimatorContract.address);
-            WaterVoucherContract.setWaterGoverningContractAddress(WaterGoverningContract.address);
+            PriceEstimatorContract = await IPriceEstimator.at(PriceEstimatorContractProxy.address);
+
+            await PriceEstimatorContract.init({
+                from: _owner
+            });
+            await PriceEstimatorContract.setHouseholdMetersContract(HouseholdMetersContract.address);
+            await PriceEstimatorContract.setWaterVouchersContract(WaterVoucherContract.address);
+
+            await WaterVoucherContract.setPriceEstimatorContractAddress(PriceEstimatorContract.address);
+            await WaterVoucherContract.setWaterGoverningContractAddress(WaterGoverningContract.address);
 
             await HouseholdMetersContract.addAdmin(_admin1);
             await HouseholdMetersContract.addHouseholdMeter(_householdMeter, _householdMembers, {
@@ -106,4 +232,9 @@ contract('PriceEstimator', function (accounts) {
         // require(_meter != address(0));
         // require(_liters > 1000);
     });
+
+    // TODO Write more test for
+    // function setHouseholdMetersContract
+    // function setWaterVouchersContract
+    // function getCurrentMonthLiters
 });