Skip to content

Commit

Permalink
Fix up tests, address comments
Browse files Browse the repository at this point in the history
  • Loading branch information
maplant committed May 16, 2024
1 parent 198a838 commit f46c818
Show file tree
Hide file tree
Showing 10 changed files with 132 additions and 139 deletions.
92 changes: 42 additions & 50 deletions mobile_verifier/src/boosting_oracles/data_sets.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ use file_store::{
};
use futures_util::{Stream, StreamExt, TryFutureExt, TryStreamExt};
use helium_proto::services::poc_mobile as proto;
use hextree::disktree::DiskTreeMap;
use lazy_static::lazy_static;
use regex::Regex;
use rust_decimal::prelude::ToPrimitive;
Expand All @@ -24,7 +23,9 @@ use tokio::{fs::File, io::AsyncWriteExt, sync::mpsc::Receiver};

use crate::{boosting_oracles::assignment::HexAssignments, coverage::SignalLevel, Settings};

use super::{footfall::Footfall, landtype::Landtype, HexAssignment, HexBoostData, Urbanization};
use super::{
footfall::Footfall, landtype::Landtype, urbanization::Urbanization, HexAssignment, HexBoostData,
};

#[async_trait::async_trait]
pub trait DataSet: HexAssignment + Send + Sync + 'static {
Expand Down Expand Up @@ -169,9 +170,7 @@ where
}
}

impl
DataSetDownloaderDaemon<Footfall<DiskTreeMap>, Landtype<DiskTreeMap>, Urbanization<DiskTreeMap>>
{
impl DataSetDownloaderDaemon<Footfall, Landtype, Urbanization> {
pub async fn create_managed_task(
pool: PgPool,
settings: &Settings,
Expand All @@ -190,9 +189,9 @@ impl
.create()
.await?;

let urbanization: Urbanization<DiskTreeMap> = Urbanization::new();
let footfall: Footfall<DiskTreeMap> = Footfall::new();
let landtype: Landtype<DiskTreeMap> = Landtype::new();
let urbanization = Urbanization::new();
let footfall = Footfall::new();
let landtype = Landtype::new();
let hex_boost_data = HexBoostData::builder()
.footfall(footfall)
.landtype(landtype)
Expand Down Expand Up @@ -548,14 +547,34 @@ pub mod db {
.fetch_one(pool)
.await?)
}

pub fn fetch_all_hexes(pool: &PgPool) -> impl Stream<Item = sqlx::Result<UnassignedHex>> + '_ {
sqlx::query_as("SELECT uuid, hex, signal_level, signal_power FROM hexes").fetch(pool)
}

pub fn fetch_hexes_with_null_assignments(
pool: &PgPool,
) -> impl Stream<Item = sqlx::Result<UnassignedHex>> + '_ {
sqlx::query_as(
"SELECT
uuid, hex, signal_level, signal_power
FROM
hexes
WHERE
urbanized IS NULL
OR footfall IS NULL
OR landtype IS NULL",
)
.fetch(pool)
}
}

pub struct AssignedCoverageObjects {
coverage_objs: HashMap<uuid::Uuid, Vec<AssignedHex>>,
pub coverage_objs: HashMap<uuid::Uuid, Vec<AssignedHex>>,
}

impl AssignedCoverageObjects {
async fn from_stream(
pub async fn assign_hex_stream(
stream: impl Stream<Item = sqlx::Result<UnassignedHex>>,
data_sets: &HexBoostData<impl HexAssignment, impl HexAssignment, impl HexAssignment>,
) -> anyhow::Result<Self> {
Expand All @@ -568,38 +587,6 @@ impl AssignedCoverageObjects {
Ok(Self { coverage_objs })
}

async fn fetch_all(
pool: &PgPool,
data_sets: &HexBoostData<impl HexAssignment, impl HexAssignment, impl HexAssignment>,
) -> anyhow::Result<Self> {
Self::from_stream(
sqlx::query_as("SELECT uuid, hex, signal_level, signal_power FROM hexes").fetch(pool),
data_sets,
)
.await
}

async fn fetch_unassigned(
pool: &PgPool,
data_sets: &HexBoostData<impl HexAssignment, impl HexAssignment, impl HexAssignment>,
) -> anyhow::Result<Self> {
Self::from_stream(
sqlx::query_as(
"SELECT
uuid, hex, signal_level, signal_power
FROM
hexes
WHERE
urbanized IS NULL
OR footfall IS NULL
OR landtype IS NULL",
)
.fetch(pool),
data_sets,
)
.await
}

async fn write(&self, boosting_reports: &FileSinkClient) -> file_store::Result {
let timestamp = Utc::now().encode_timestamp();
for (uuid, hexes) in self.coverage_objs.iter() {
Expand Down Expand Up @@ -635,7 +622,7 @@ impl AssignedCoverageObjects {
Ok(())
}

async fn save(self, pool: &PgPool) -> anyhow::Result<()> {
pub async fn save(self, pool: &PgPool) -> anyhow::Result<()> {
const NUMBER_OF_FIELDS_IN_QUERY: u16 = 7;
const ASSIGNMENTS_MAX_BATCH_ENTRIES: usize =
(u16::MAX / NUMBER_OF_FIELDS_IN_QUERY) as usize;
Expand Down Expand Up @@ -702,19 +689,20 @@ impl UnassignedHex {
}

pub struct AssignedHex {
uuid: uuid::Uuid,
hex: u64,
signal_level: SignalLevel,
signal_power: i32,
assignments: HexAssignments,
pub uuid: uuid::Uuid,
pub hex: u64,
pub signal_level: SignalLevel,
pub signal_power: i32,
pub assignments: HexAssignments,
}

pub async fn set_all_oracle_boosting_assignments(
pool: &PgPool,
data_sets: &HexBoostData<impl HexAssignment, impl HexAssignment, impl HexAssignment>,
file_sink: &FileSinkClient,
) -> anyhow::Result<()> {
let assigned_coverage_objs = AssignedCoverageObjects::fetch_all(pool, data_sets).await?;
let assigned_coverage_objs =
AssignedCoverageObjects::assign_hex_stream(db::fetch_all_hexes(pool), data_sets).await?;
assigned_coverage_objs.write(file_sink).await?;
assigned_coverage_objs.save(pool).await?;
Ok(())
Expand All @@ -725,7 +713,11 @@ pub async fn set_unassigned_oracle_boosting_assignments(
data_sets: &HexBoostData<impl HexAssignment, impl HexAssignment, impl HexAssignment>,
file_sink: &FileSinkClient,
) -> anyhow::Result<()> {
let assigned_coverage_objs = AssignedCoverageObjects::fetch_unassigned(pool, data_sets).await?;
let assigned_coverage_objs = AssignedCoverageObjects::assign_hex_stream(
db::fetch_hexes_with_null_assignments(pool),
data_sets,
)
.await?;
assigned_coverage_objs.write(file_sink).await?;
assigned_coverage_objs.save(pool).await?;
Ok(())
Expand Down
19 changes: 8 additions & 11 deletions mobile_verifier/src/boosting_oracles/footfall.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,37 +3,37 @@ use std::path::Path;
use chrono::{DateTime, Utc};
use hextree::disktree::DiskTreeMap;

use super::{Assignment, DataSet, DataSetType, DiskTreeLike, HexAssignment};
use super::{Assignment, DataSet, DataSetType, HexAssignment};

pub struct Footfall<Foot> {
footfall: Option<Foot>,
pub struct Footfall {
footfall: Option<DiskTreeMap>,
timestamp: Option<DateTime<Utc>>,
}

impl<F> Footfall<F> {
impl Footfall {
pub fn new() -> Self {
Self {
footfall: None,
timestamp: None,
}
}

pub fn new_mock(footfall: F) -> Self {
pub fn new_mock(footfall: DiskTreeMap) -> Self {
Self {
footfall: Some(footfall),
timestamp: None,
}
}
}

impl<F> Default for Footfall<F> {
impl Default for Footfall {
fn default() -> Self {
Self::new()
}
}

#[async_trait::async_trait]
impl DataSet for Footfall<DiskTreeMap> {
impl DataSet for Footfall {
const TYPE: DataSetType = DataSetType::Footfall;

fn timestamp(&self) -> Option<DateTime<Utc>> {
Expand All @@ -51,10 +51,7 @@ impl DataSet for Footfall<DiskTreeMap> {
}
}

impl<Foot> HexAssignment for Footfall<Foot>
where
Foot: DiskTreeLike + Send + Sync + 'static,
{
impl HexAssignment for Footfall {
fn assignment(&self, cell: hextree::Cell) -> anyhow::Result<Assignment> {
let Some(ref footfall) = self.footfall else {
anyhow::bail!("No footfall data set has been loaded");
Expand Down
19 changes: 8 additions & 11 deletions mobile_verifier/src/boosting_oracles/landtype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,37 +3,37 @@ use std::path::Path;
use chrono::{DateTime, Utc};
use hextree::disktree::DiskTreeMap;

use super::{Assignment, DataSet, DataSetType, DiskTreeLike, HexAssignment};
use super::{Assignment, DataSet, DataSetType, HexAssignment};

pub struct Landtype<Land> {
landtype: Option<Land>,
pub struct Landtype {
landtype: Option<DiskTreeMap>,
timestamp: Option<DateTime<Utc>>,
}

impl<L> Landtype<L> {
impl Landtype {
pub fn new() -> Self {
Self {
landtype: None,
timestamp: None,
}
}

pub fn new_mock(landtype: L) -> Self {
pub fn new_mock(landtype: DiskTreeMap) -> Self {
Self {
landtype: Some(landtype),
timestamp: None,
}
}
}

impl<L> Default for Landtype<L> {
impl Default for Landtype {
fn default() -> Self {
Self::new()
}
}

#[async_trait::async_trait]
impl DataSet for Landtype<DiskTreeMap> {
impl DataSet for Landtype {
const TYPE: DataSetType = DataSetType::Landtype;

fn timestamp(&self) -> Option<DateTime<Utc>> {
Expand Down Expand Up @@ -131,10 +131,7 @@ impl From<LandtypeValue> for Assignment {
}
}

impl<Land> HexAssignment for Landtype<Land>
where
Land: DiskTreeLike + Send + Sync + 'static,
{
impl HexAssignment for Landtype {
fn assignment(&self, cell: hextree::Cell) -> anyhow::Result<Assignment> {
let Some(ref landtype) = self.landtype else {
anyhow::bail!("No landtype data set has been loaded");
Expand Down
31 changes: 2 additions & 29 deletions mobile_verifier/src/boosting_oracles/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,6 @@ use crate::boosting_oracles::assignment::HexAssignments;
pub use assignment::Assignment;
pub use data_sets::*;

use hextree::disktree::DiskTreeMap;
pub use urbanization::Urbanization;

pub trait HexAssignment: Send + Sync + 'static {
fn assignment(&self, cell: hextree::Cell) -> anyhow::Result<Assignment>;
}
Expand Down Expand Up @@ -68,38 +65,14 @@ where
}
}

trait DiskTreeLike: Send + Sync {
fn get(&self, cell: hextree::Cell) -> hextree::Result<Option<(hextree::Cell, &[u8])>>;
}

impl DiskTreeLike for DiskTreeMap {
fn get(&self, cell: hextree::Cell) -> hextree::Result<Option<(hextree::Cell, &[u8])>> {
self.get(cell)
}
}

impl DiskTreeLike for std::collections::HashSet<hextree::Cell> {
fn get(&self, cell: hextree::Cell) -> hextree::Result<Option<(hextree::Cell, &[u8])>> {
Ok(self.contains(&cell).then_some((cell, &[])))
}
}

pub struct MockDiskTree;

impl DiskTreeLike for MockDiskTree {
fn get(&self, cell: hextree::Cell) -> hextree::Result<Option<(hextree::Cell, &[u8])>> {
Ok(Some((cell, &[])))
}
}

#[cfg(test)]
mod tests {

use std::io::Cursor;

use hextree::HexTreeMap;
use hextree::{disktree::DiskTreeMap, HexTreeMap};

use self::{footfall::Footfall, landtype::Landtype};
use self::{footfall::Footfall, landtype::Landtype, urbanization::Urbanization};

use super::*;

Expand Down
19 changes: 8 additions & 11 deletions mobile_verifier/src/boosting_oracles/urbanization.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,37 +3,37 @@ use std::path::Path;
use chrono::{DateTime, Utc};
use hextree::disktree::DiskTreeMap;

use super::{Assignment, DataSet, DataSetType, DiskTreeLike, HexAssignment};
use super::{Assignment, DataSet, DataSetType, HexAssignment};

pub struct Urbanization<DT> {
urbanized: Option<DT>,
pub struct Urbanization {
urbanized: Option<DiskTreeMap>,
timestamp: Option<DateTime<Utc>>,
}

impl<DT> Urbanization<DT> {
impl Urbanization {
pub fn new() -> Self {
Self {
urbanized: None,
timestamp: None,
}
}

pub fn new_mock(urbanized: DT) -> Self {
pub fn new_mock(urbanized: DiskTreeMap) -> Self {
Self {
urbanized: Some(urbanized),
timestamp: None,
}
}
}

impl<DT> Default for Urbanization<DT> {
impl Default for Urbanization {
fn default() -> Self {
Self::new()
}
}

#[async_trait::async_trait]
impl DataSet for Urbanization<DiskTreeMap> {
impl DataSet for Urbanization {
const TYPE: DataSetType = DataSetType::Urbanization;

fn timestamp(&self) -> Option<DateTime<Utc>> {
Expand All @@ -51,10 +51,7 @@ impl DataSet for Urbanization<DiskTreeMap> {
}
}

impl<Urban> HexAssignment for Urbanization<Urban>
where
Urban: DiskTreeLike + Send + Sync + 'static,
{
impl HexAssignment for Urbanization {
fn assignment(&self, cell: hextree::Cell) -> anyhow::Result<Assignment> {
let Some(ref urbanized) = self.urbanized else {
anyhow::bail!("No urbanization data set has been loaded");
Expand Down
Loading

0 comments on commit f46c818

Please sign in to comment.