test refactor and fix

open-web3-stack · Oct 22, 2024 · 544c6cf · 544c6cf
1 parent ab09963
commit 544c6cf
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 50 deletions.
diff --git a/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift b/Blockchain/Tests/BlockchainTests/BlockAuthorTests.swift
@@ -6,64 +6,53 @@ import Utils
 @testable import Blockchain
 
 struct BlockAuthorTests {
-    let config: ProtocolConfigRef
-    let timeProvider: MockTimeProvider
-    let dataProvider: BlockchainDataProvider
-    let eventBus: EventBus
-    let scheduler: MockScheduler
-    let keystore: KeyStore
-    let blockAuthor: BlockAuthor
-    let runtime: Runtime
-    let storeMiddleware: StoreMiddleware
+    func setup() async -> (BlockchainServices, BlockAuthor, Runtime) {
+        setupTestLogger()
 
-    init() async throws {
-        // setupTestLogger()
-
-        config = ProtocolConfigRef.dev
-        timeProvider = MockTimeProvider(time: 988)
-
-        let (genesisState, genesisBlock) = try State.devGenesis(config: config)
-        dataProvider = try await BlockchainDataProvider(InMemoryDataProvider(genesisState: genesisState, genesisBlock: genesisBlock))
-
-        storeMiddleware = StoreMiddleware()
-        eventBus = EventBus(eventMiddleware: .serial(Middleware(storeMiddleware), .noError), handlerMiddleware: .noError)
-
-        scheduler = MockScheduler(timeProvider: timeProvider)
-
-        keystore = try await DevKeyStore(devKeysCount: config.value.totalNumberOfValidators)
-
-        blockAuthor = await BlockAuthor(
-            config: config,
-            dataProvider: dataProvider,
-            eventBus: eventBus,
-            keystore: keystore,
-            scheduler: scheduler,
-            extrinsicPool: ExtrinsicPoolService(config: config, dataProvider: dataProvider, eventBus: eventBus)
-        )
-
-        runtime = Runtime(config: config)
+        let services = await BlockchainServices()
+        let blockAuthor = await services.blockAuthor()
+        let runtime = Runtime(config: services.config)
+        return (services, blockAuthor, runtime)
     }
 
     @Test
     func createNewBlockWithFallbackKey() async throws {
-        let genesisState = try await dataProvider.getState(hash: dataProvider.genesisBlockHash)
+        let (services, blockAuthor, runtime) = await setup()
+        let config = services.config
+        let timeProvider = services.timeProvider
+        let genesisState = services.genesisState
 
         let timeslot = timeProvider.getTime().timeToTimeslot(config: config)
 
-        // get the validator key
-        let idx = timeslot % UInt32(config.value.totalNumberOfValidators)
-        let devKey = try DevKeyStore.getDevKey(seed: idx)
+        // dry run Safrole to get the validator key
+        let res = try genesisState.value.updateSafrole(
+            config: config,
+            slot: timeslot,
+            entropy: Data32(),
+            offenders: [],
+            extrinsics: .dummy(config: config)
+        )
+
+        let idx = timeslot % UInt32(config.value.epochLength)
+        let key = res.state.ticketsOrKeys.right!.array[Int(idx)]
+        let pubkey = try! Bandersnatch.PublicKey(data: key)
 
         // Create a new block
-        let block = try await blockAuthor.createNewBlock(timeslot: timeslot, claim: .right(devKey.bandersnatch))
+        let block = try await blockAuthor.createNewBlock(timeslot: timeslot, claim: .right(pubkey))
 
         // Verify block
         try _ = runtime.apply(block: block, state: genesisState, context: .init(timeslot: timeslot + 1))
     }
 
     @Test
     func createNewBlockWithTicket() async throws {
-        let genesisState = try await dataProvider.getState(hash: dataProvider.genesisBlockHash)
+        let (services, blockAuthor, runtime) = await setup()
+        let config = services.config
+        let timeProvider = services.timeProvider
+        let genesisState = services.genesisState
+        let keystore = services.keystore
+        let dataProvider = services.dataProvider
+
         var state = genesisState.value
 
         state.safroleState.ticketsVerifier = try Bandersnatch.RingCommitment(
@@ -109,12 +98,19 @@ struct BlockAuthorTests {
 
     @Test
     func firstBlock() async throws {
-        let genesisState = try await dataProvider.getState(hash: dataProvider.genesisBlockHash)
+        let (services, blockAuthor, runtime) = await setup()
+        let config = services.config
+        let timeProvider = services.timeProvider
+        let genesisState = services.genesisState
+        let scheduler = services.scheduler
+        let storeMiddleware = services.storeMiddleware
 
         await blockAuthor.on(genesis: genesisState)
 
         #expect(scheduler.storage.value.tasks.count > 0)
 
+        await scheduler.advance(by: TimeInterval(2))
+
         let events = await storeMiddleware.wait()
         #expect(events.count == 1)
         #expect(events.first is RuntimeEvents.BlockAuthored)

diff --git a/Blockchain/Tests/BlockchainTests/BlockchainServices.swift b/Blockchain/Tests/BlockchainTests/BlockchainServices.swift
@@ -9,34 +9,46 @@ struct BlockchainServices {
     let scheduler: MockScheduler
     let keystore: KeyStore
     let storeMiddleware: StoreMiddleware
-    let blockchain: Blockchain
     let genesisBlock: BlockRef
     let genesisState: StateRef
 
     init(
         config: ProtocolConfigRef = .dev,
         timeProvider: MockTimeProvider = MockTimeProvider(time: 988)
-    ) async throws {
+    ) async {
         self.config = config
         self.timeProvider = timeProvider
 
-        let (genesisState, genesisBlock) = try State.devGenesis(config: config)
+        let (genesisState, genesisBlock) = try! State.devGenesis(config: config)
         self.genesisBlock = genesisBlock
         self.genesisState = genesisState
-        dataProvider = try await BlockchainDataProvider(InMemoryDataProvider(genesisState: genesisState, genesisBlock: genesisBlock))
+        dataProvider = try! await BlockchainDataProvider(InMemoryDataProvider(genesisState: genesisState, genesisBlock: genesisBlock))
 
         storeMiddleware = StoreMiddleware()
         eventBus = EventBus(eventMiddleware: .serial(Middleware(storeMiddleware), .noError), handlerMiddleware: .noError)
 
         scheduler = MockScheduler(timeProvider: timeProvider)
 
-        keystore = try await DevKeyStore()
+        keystore = try! await DevKeyStore()
+    }
 
-        blockchain = try await Blockchain(
+    func blockchain() async -> Blockchain {
+        try! await Blockchain(
             config: config,
             dataProvider: dataProvider,
             timeProvider: timeProvider,
             eventBus: eventBus
         )
     }
+
+    func blockAuthor() async -> BlockAuthor {
+        await BlockAuthor(
+            config: config,
+            dataProvider: dataProvider,
+            eventBus: eventBus,
+            keystore: keystore,
+            scheduler: scheduler,
+            extrinsicPool: ExtrinsicPoolService(config: config, dataProvider: dataProvider, eventBus: eventBus)
+        )
+    }
 }
diff --git a/Blockchain/Tests/BlockchainTests/ValidatorServiceTests.swift b/Blockchain/Tests/BlockchainTests/ValidatorServiceTests.swift
@@ -9,11 +9,11 @@ struct ValidatorServiceTests {
     func setup(time: TimeInterval = 988) async throws -> (BlockchainServices, ValidatorService) {
         // setupTestLogger()
 
-        let services = try await BlockchainServices(
+        let services = await BlockchainServices(
             timeProvider: MockTimeProvider(time: time)
         )
         let validatorService = await ValidatorService(
-            blockchain: services.blockchain,
+            blockchain: services.blockchain(),
             keystore: services.keystore,
             eventBus: services.eventBus,
             scheduler: services.scheduler,